pinecone-research

1"""Pinecone memory manager — namespace-based session memory for agents.2 3Provides helpers for storing and retrieving agent conversation memory4using Pinecone namespaces. Each session gets its own namespace for isolation,5with cross-session search available via the global namespace.6 7Usage:8    export PINECONE_API_KEY="your-key"9    export OPENAI_API_KEY="your-key"10    python memory_manager.py --index-name agent-memory --action store \11        --session-id sess-001 --text "User discussed project architecture"12    python memory_manager.py --index-name agent-memory --action recall \13        --query "architecture decisions"14    python memory_manager.py --index-name agent-memory --action cleanup \15        --session-id sess-00116"""17from __future__ import annotations18 19import argparse20import hashlib21import os22import sys23import time24 25 26def get_pinecone_client():27    """Initialize Pinecone client from environment."""28    try:29        from pinecone import Pinecone30    except ImportError:31        print("Error: pinecone-client not installed. Run: pip install pinecone-client", file=sys.stderr)32        sys.exit(1)33 34    api_key = os.environ.get("PINECONE_API_KEY")35    if not api_key:36        print("Error: PINECONE_API_KEY environment variable not set.", file=sys.stderr)37        sys.exit(1)38 39    return Pinecone(api_key=api_key)40 41 42def get_embeddings():43    """Get the embedding model."""44    try:45        from langchain_openai import OpenAIEmbeddings46    except ImportError:47        print("Error: langchain-openai not installed. Run: pip install langchain-openai", file=sys.stderr)48        sys.exit(1)49    return OpenAIEmbeddings()50 51 52def store_memory(index, session_id: str, text: str, metadata: dict | None = None):53    """Store a memory entry in the session namespace."""54    embeddings = get_embeddings()55    vector = embeddings.embed_query(text)56 57    doc_id = hashlib.sha256(f"{session_id}:{text}:{time.time()}".encode()).hexdigest()[:16]58    entry_metadata = {59        "text": text[:1000],60        "session_id": session_id,61        "timestamp": int(time.time()),62    }63    if metadata:64        entry_metadata.update(metadata)65 66    index.upsert(67        vectors=[{"id": doc_id, "values": vector, "metadata": entry_metadata}],68        namespace=session_id,69    )70    print(f"Stored memory [{doc_id}] in namespace '{session_id}'")71    return doc_id72 73 74def recall_memories(index, query: str, session_id: str | None = None, top_k: int = 5):75    """Recall memories matching a query, optionally scoped to a session."""76    embeddings = get_embeddings()77    query_vector = embeddings.embed_query(query)78 79    kwargs = {"vector": query_vector, "top_k": top_k, "include_metadata": True}80    if session_id:81        kwargs["namespace"] = session_id82 83    results = index.query(**kwargs)84 85    print(f"\nRecalling memories for: {query!r}")86    if session_id:87        print(f"Scoped to session: {session_id}")88    print(f"Found {len(results['matches'])} results:\n")89 90    for match in results["matches"]:91        score = match["score"]92        text = match["metadata"].get("text", "")[:200]93        sess = match["metadata"].get("session_id", "unknown")94        ts = match["metadata"].get("timestamp", 0)95        print(f"  [{score:.4f}] session={sess} time={ts}")96        print(f"    {text}")97        print()98 99    return results100 101 102def cleanup_session(index, session_id: str):103    """Delete all vectors in a session namespace."""104    index.delete(delete_all=True, namespace=session_id)105    print(f"Cleaned up namespace '{session_id}'")106 107 108def show_stats(index):109    """Show index statistics."""110    stats = index.describe_index_stats()111    print(f"Total vectors: {stats['total_vector_count']}")112    namespaces = stats.get("namespaces", {})113    if namespaces:114        print(f"Namespaces ({len(namespaces)}):")115        for ns, info in sorted(namespaces.items()):116            print(f"  '{ns}': {info['vector_count']} vectors")117    else:118        print("No namespaces found.")119 120 121def main():122    parser = argparse.ArgumentParser(description="Pinecone agent memory manager")123    parser.add_argument("--index-name", required=True, help="Pinecone index name")124    parser.add_argument(125        "--action",126        choices=["store", "recall", "cleanup", "stats"],127        required=True,128    )129    parser.add_argument("--session-id", help="Session namespace ID")130    parser.add_argument("--text", help="Text to store as memory")131    parser.add_argument("--query", help="Query for recall")132    parser.add_argument("--top-k", type=int, default=5, help="Number of results")133    args = parser.parse_args()134 135    pc = get_pinecone_client()136    index = pc.Index(args.index_name)137 138    if args.action == "store":139        if not args.session_id or not args.text:140            parser.error("--session-id and --text required for store action")141        store_memory(index, args.session_id, args.text)142    elif args.action == "recall":143        if not args.query:144            parser.error("--query required for recall action")145        recall_memories(index, args.query, session_id=args.session_id, top_k=args.top_k)146    elif args.action == "cleanup":147        if not args.session_id:148            parser.error("--session-id required for cleanup action")149        cleanup_session(index, args.session_id)150    elif args.action == "stats":151        show_stats(index)152 153 154if __name__ == "__main__":155    main()156