pinecone-research

1"""Pinecone RAG pipeline — index documents and query with retrieval-augmented generation.2 3Usage:4    export PINECONE_API_KEY="your-key"5    export OPENAI_API_KEY="your-key"6    python rag_pipeline.py --index-name agent-memory --action index --docs-dir ./docs7    python rag_pipeline.py --index-name agent-memory --action query --query "How does X work?"8"""9from __future__ import annotations10 11import argparse12import os13import sys14from pathlib import Path15 16 17def get_pinecone_client():18    """Initialize Pinecone client from environment."""19    try:20        from pinecone import Pinecone21    except ImportError:22        print("Error: pinecone-client not installed. Run: pip install pinecone-client", file=sys.stderr)23        sys.exit(1)24 25    api_key = os.environ.get("PINECONE_API_KEY")26    if not api_key:27        print("Error: PINECONE_API_KEY environment variable not set.", file=sys.stderr)28        sys.exit(1)29 30    return Pinecone(api_key=api_key)31 32 33def ensure_index(pc, index_name: str, dimension: int = 1536):34    """Create the index if it doesn't exist."""35    from pinecone import ServerlessSpec36 37    existing = [idx.name for idx in pc.list_indexes()]38    if index_name not in existing:39        pc.create_index(40            name=index_name,41            dimension=dimension,42            metric="cosine",43            spec=ServerlessSpec(cloud="aws", region="us-east-1"),44        )45        print(f"Created index: {index_name}")46    else:47        print(f"Index already exists: {index_name}")48    return pc.Index(index_name)49 50 51def load_documents(docs_dir: str) -> list[dict]:52    """Load text files from a directory as documents."""53    docs = []54    docs_path = Path(docs_dir)55    if not docs_path.is_dir():56        print(f"Error: {docs_dir} is not a directory.", file=sys.stderr)57        sys.exit(1)58 59    for filepath in sorted(docs_path.rglob("*.txt")):60        text = filepath.read_text(encoding="utf-8").strip()61        if text:62            docs.append({63                "id": str(filepath.relative_to(docs_path)),64                "text": text,65                "metadata": {"source": str(filepath.name)},66            })67    return docs68 69 70def index_documents(index, docs: list[dict], batch_size: int = 100):71    """Embed and upsert documents into Pinecone."""72    try:73        from langchain_openai import OpenAIEmbeddings74    except ImportError:75        print("Error: langchain-openai not installed. Run: pip install langchain-openai", file=sys.stderr)76        sys.exit(1)77 78    embeddings = OpenAIEmbeddings()79    vectors = []80 81    for doc in docs:82        embedding = embeddings.embed_query(doc["text"])83        vectors.append({84            "id": doc["id"],85            "values": embedding,86            "metadata": {**doc["metadata"], "text": doc["text"][:1000]},87        })88 89    # Batch upsert90    for i in range(0, len(vectors), batch_size):91        batch = vectors[i : i + batch_size]92        index.upsert(vectors=batch)93        print(f"Upserted batch {i // batch_size + 1} ({len(batch)} vectors)")94 95    print(f"Total vectors indexed: {len(vectors)}")96 97 98def query_index(index, query: str, top_k: int = 5):99    """Embed a query and retrieve similar documents from Pinecone."""100    try:101        from langchain_openai import OpenAIEmbeddings102    except ImportError:103        print("Error: langchain-openai not installed. Run: pip install langchain-openai", file=sys.stderr)104        sys.exit(1)105 106    embeddings = OpenAIEmbeddings()107    query_vector = embeddings.embed_query(query)108 109    results = index.query(vector=query_vector, top_k=top_k, include_metadata=True)110 111    print(f"\nQuery: {query}")112    print(f"Top {top_k} results:\n")113    for match in results["matches"]:114        score = match["score"]115        source = match["metadata"].get("source", "unknown")116        text_preview = match["metadata"].get("text", "")[:200]117        print(f"  [{score:.4f}] {source}")118        print(f"    {text_preview}...")119        print()120 121    return results122 123 124def main():125    parser = argparse.ArgumentParser(description="Pinecone RAG pipeline")126    parser.add_argument("--index-name", required=True, help="Pinecone index name")127    parser.add_argument("--action", choices=["index", "query", "stats"], required=True)128    parser.add_argument("--docs-dir", help="Directory of .txt files to index")129    parser.add_argument("--query", help="Query string for retrieval")130    parser.add_argument("--top-k", type=int, default=5, help="Number of results to return")131    args = parser.parse_args()132 133    pc = get_pinecone_client()134    index = ensure_index(pc, args.index_name)135 136    if args.action == "index":137        if not args.docs_dir:138            parser.error("--docs-dir required for index action")139        docs = load_documents(args.docs_dir)140        if not docs:141            print("No .txt documents found.", file=sys.stderr)142            sys.exit(1)143        index_documents(index, docs)144    elif args.action == "query":145        if not args.query:146            parser.error("--query required for query action")147        query_index(index, args.query, top_k=args.top_k)148    elif args.action == "stats":149        stats = index.describe_index_stats()150        print(f"Total vectors: {stats['total_vector_count']}")151        for ns, info in stats.get("namespaces", {}).items():152            print(f"  Namespace '{ns}': {info['vector_count']} vectors")153 154 155if __name__ == "__main__":156    main()157