157 lines
5.2 KiB
Python
157 lines
5.2 KiB
Python
"""Pinecone RAG pipeline — index documents and query with retrieval-augmented generation.
|
|
|
|
Usage:
|
|
export PINECONE_API_KEY="your-key"
|
|
export OPENAI_API_KEY="your-key"
|
|
python rag_pipeline.py --index-name agent-memory --action index --docs-dir ./docs
|
|
python rag_pipeline.py --index-name agent-memory --action query --query "How does X work?"
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def get_pinecone_client():
|
|
"""Initialize Pinecone client from environment."""
|
|
try:
|
|
from pinecone import Pinecone
|
|
except ImportError:
|
|
print("Error: pinecone-client not installed. Run: pip install pinecone-client", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
api_key = os.environ.get("PINECONE_API_KEY")
|
|
if not api_key:
|
|
print("Error: PINECONE_API_KEY environment variable not set.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
return Pinecone(api_key=api_key)
|
|
|
|
|
|
def ensure_index(pc, index_name: str, dimension: int = 1536):
|
|
"""Create the index if it doesn't exist."""
|
|
from pinecone import ServerlessSpec
|
|
|
|
existing = [idx.name for idx in pc.list_indexes()]
|
|
if index_name not in existing:
|
|
pc.create_index(
|
|
name=index_name,
|
|
dimension=dimension,
|
|
metric="cosine",
|
|
spec=ServerlessSpec(cloud="aws", region="us-east-1"),
|
|
)
|
|
print(f"Created index: {index_name}")
|
|
else:
|
|
print(f"Index already exists: {index_name}")
|
|
return pc.Index(index_name)
|
|
|
|
|
|
def load_documents(docs_dir: str) -> list[dict]:
|
|
"""Load text files from a directory as documents."""
|
|
docs = []
|
|
docs_path = Path(docs_dir)
|
|
if not docs_path.is_dir():
|
|
print(f"Error: {docs_dir} is not a directory.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
for filepath in sorted(docs_path.rglob("*.txt")):
|
|
text = filepath.read_text(encoding="utf-8").strip()
|
|
if text:
|
|
docs.append({
|
|
"id": str(filepath.relative_to(docs_path)),
|
|
"text": text,
|
|
"metadata": {"source": str(filepath.name)},
|
|
})
|
|
return docs
|
|
|
|
|
|
def index_documents(index, docs: list[dict], batch_size: int = 100):
|
|
"""Embed and upsert documents into Pinecone."""
|
|
try:
|
|
from langchain_openai import OpenAIEmbeddings
|
|
except ImportError:
|
|
print("Error: langchain-openai not installed. Run: pip install langchain-openai", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
embeddings = OpenAIEmbeddings()
|
|
vectors = []
|
|
|
|
for doc in docs:
|
|
embedding = embeddings.embed_query(doc["text"])
|
|
vectors.append({
|
|
"id": doc["id"],
|
|
"values": embedding,
|
|
"metadata": {**doc["metadata"], "text": doc["text"][:1000]},
|
|
})
|
|
|
|
# Batch upsert
|
|
for i in range(0, len(vectors), batch_size):
|
|
batch = vectors[i : i + batch_size]
|
|
index.upsert(vectors=batch)
|
|
print(f"Upserted batch {i // batch_size + 1} ({len(batch)} vectors)")
|
|
|
|
print(f"Total vectors indexed: {len(vectors)}")
|
|
|
|
|
|
def query_index(index, query: str, top_k: int = 5):
|
|
"""Embed a query and retrieve similar documents from Pinecone."""
|
|
try:
|
|
from langchain_openai import OpenAIEmbeddings
|
|
except ImportError:
|
|
print("Error: langchain-openai not installed. Run: pip install langchain-openai", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
embeddings = OpenAIEmbeddings()
|
|
query_vector = embeddings.embed_query(query)
|
|
|
|
results = index.query(vector=query_vector, top_k=top_k, include_metadata=True)
|
|
|
|
print(f"\nQuery: {query}")
|
|
print(f"Top {top_k} results:\n")
|
|
for match in results["matches"]:
|
|
score = match["score"]
|
|
source = match["metadata"].get("source", "unknown")
|
|
text_preview = match["metadata"].get("text", "")[:200]
|
|
print(f" [{score:.4f}] {source}")
|
|
print(f" {text_preview}...")
|
|
print()
|
|
|
|
return results
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Pinecone RAG pipeline")
|
|
parser.add_argument("--index-name", required=True, help="Pinecone index name")
|
|
parser.add_argument("--action", choices=["index", "query", "stats"], required=True)
|
|
parser.add_argument("--docs-dir", help="Directory of .txt files to index")
|
|
parser.add_argument("--query", help="Query string for retrieval")
|
|
parser.add_argument("--top-k", type=int, default=5, help="Number of results to return")
|
|
args = parser.parse_args()
|
|
|
|
pc = get_pinecone_client()
|
|
index = ensure_index(pc, args.index_name)
|
|
|
|
if args.action == "index":
|
|
if not args.docs_dir:
|
|
parser.error("--docs-dir required for index action")
|
|
docs = load_documents(args.docs_dir)
|
|
if not docs:
|
|
print("No .txt documents found.", file=sys.stderr)
|
|
sys.exit(1)
|
|
index_documents(index, docs)
|
|
elif args.action == "query":
|
|
if not args.query:
|
|
parser.error("--query required for query action")
|
|
query_index(index, args.query, top_k=args.top_k)
|
|
elif args.action == "stats":
|
|
stats = index.describe_index_stats()
|
|
print(f"Total vectors: {stats['total_vector_count']}")
|
|
for ns, info in stats.get("namespaces", {}).items():
|
|
print(f" Namespace '{ns}': {info['vector_count']} vectors")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|