From 8898d37185c4c88fab88cfb01d9d2b7185b37518 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 13:37:46 +0000 Subject: [PATCH] add agent.py --- agent.py | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 agent.py diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..15fdf09 --- /dev/null +++ b/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