from pinecone import Pinecone, ServerlessSpec
from openai import OpenAI
import os

# Setup
pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
openai_client = OpenAI()

# Create index
if "knowledge-base" not in pc.list_indexes().names():
    pc.create_index(
        name="knowledge-base",
        dimension=1536,  # text-embedding-3-small dimensions
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
    )

index = pc.Index("knowledge-base")

def embed(texts: list[str]) -> list[list[float]]:
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]

# Upsert documents (batch for efficiency)
def index_documents(docs: list[dict]):
    texts = [d["text"] for d in docs]
    embeddings = embed(texts)

    vectors = [
        {
            "id": doc["id"],
            "values": emb,
            "metadata": {"text": doc["text"], "source": doc.get("source", "")}
        }
        for doc, emb in zip(docs, embeddings)
    ]

    # Batch upsert for efficiency
    batch_size = 100
    for i in range(0, len(vectors), batch_size):
        index.upsert(vectors=vectors[i:i+batch_size])

# Query
def search(query: str, top_k: int = 5, filter: dict = None) -> list[dict]:
    query_embedding = embed([query])[0]

    results = index.query(
        vector=query_embedding,
        top_k=top_k,
        include_metadata=True,
        filter=filter  # e.g., {"source": {"$eq": "legal"}}
    )

    return [
        {"text": r.metadata["text"], "score": r.score, "id": r.id}
        for r in results.matches
    ]
Vector Databases — Embeddings, Similarity Search

Vector Databases — Embeddings, Similarity Search

Vector Databases — Embeddings, Similarity Search
