add agent.py

agent.py

@@ -0,0 +1,78 @@
+"""RAG agent: Qdrant vector store + Ollama LLM + LangGraph."""
+import os
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain_core.documents import Document
+from langchain.tools import tool
+from langgraph.prebuilt import create_react_agent
+
+QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
+COLLECTION_NAME = "knowledge"
+
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
+_vs = None
+
+
+def get_vector_store() -> QdrantVectorStore:
+    """Return singleton Qdrant vector store, creating collection if needed."""
+    global _vs
+    if _vs is None:
+        from qdrant_client import QdrantClient
+        from qdrant_client.models import Distance, VectorParams
+        client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+        existing = [c.name for c in client.get_collections().collections]
+        if COLLECTION_NAME not in existing:
+            client.create_collection(
+                collection_name=COLLECTION_NAME,
+                vectors_config=VectorParams(size=768, distance=Distance.COSINE),
+            )
+        _vs = QdrantVectorStore(
+            client=client,
+            collection_name=COLLECTION_NAME,
+            embedding=embeddings,
+        )
+    return _vs
+
+
+@tool
+def search_knowledge_base(query: str) -> str:
+    """Search the knowledge base for relevant passages.
+
+    Args:
+        query: search query string
+    """
+    docs = get_vector_store().similarity_search(query, k=5)
+    if not docs:
+        return "No results found."
+    return "\n\n".join(f"{i+1}. {d.page_content}" for i, d in enumerate(docs))
+
+
+@tool
+def add_to_knowledge_base(text: str) -> str:
+    """Add a text passage to the knowledge base.
+
+    Args:
+        text: text to store
+    """
+    get_vector_store().add_documents([Document(page_content=text)])
+    return "Added to knowledge base."
+
+llm = ChatOllama(model="llama3", temperature=0.0)
+
+agent = create_react_agent(
+    model=llm,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    state_modifier=(
+        "You are a helpful assistant with access to a local knowledge base. "
+        "Use search_knowledge_base to find information. "
+        "Use add_to_knowledge_base to store new information."
+    ),
+)
+
+
+def run_agent(user_input: str) -> str:
+    """Run the agent and return the last message content."""
+    from langchain_core.messages import HumanMessage
+    result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
+    return result["messages"][-1].content