add agent.py

agent.py

@@ -1,40 +1,79 @@
 """
-Agent creation using LangChain create_agent.
+Agent definition for the RAG system.
+
+Provides two tools:
+* search_knowledge_base(query, max_results)
+* add_to_knowledge_base(content, title)
+
+The agent is created with create_agent from langchain.agents.
 """
+
 import os
-from langchain_ollama import OllamaLLM
+from typing import List
-from langchain_core.messages import HumanMessage
+from langchain_ollama import OllamaEmbeddings
-from langchain.agents import create_agent
+from chromadb import PersistentClient
+from chromadb.utils import embedding_functions as ef
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from tools import search_knowledge_base, add_to_knowledge_base
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage
 
-# LLM via Ollama
+# Initialize embeddings and Chroma client
-llm = OllamaLLM(model="llama3")
+EMBEDDINGS = OllamaEmbeddings(model="nomic-embed-text")
+CHROMA_PATH = os.path.join(os.getcwd(), "chromadb_store")
+CLIENT = PersistentClient(path=CHROMA_PATH)
+COLLECTION_NAME = "knowledge"
+if COLLECTION_NAME not in CLIENT.list_collections():
+    CLIENT.create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
+COLL = CLIENT.get_or_create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS)
 
-# System prompt instructing to use knowledge base tools
+# Chunker from chunker.py
-SYSTEM_PROMPT = """
+from chunker import CHUNKER
-You are a helpful assistant that can store and retrieve information.
-Use the provided tools search_knowledge_base and add_to_knowledge_base.
-When answering, prefer to call the tools if needed.
-"""
 
-def create_rag_agent():
+@tool
-    agent = create_agent(
+def add_to_knowledge_base(content: str, title: str) -> str:
-        llm=llm,
+    """
-        tools=[search_knowledge_base, add_to_knowledge_base],
+    Add a document to the knowledge base.
-        system_prompt=SYSTEM_PROMPT,
+    The content is split into chunks and stored with metadata.
+    Returns confirmation message.
+    """
+    # Split content
+    chunks = CHUNKER.split_text(content)
+    ids = [f"{title}_{i}" for i in range(len(chunks))]
+    metadatas = [{"title": title} for _ in chunks]
+    COLL.add(ids=ids, documents=chunks, metadatas=metadatas)
+    return f"Added {len(chunks)} chunks from '{title}'."
+
+@tool
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """
+    Search the knowledge base for relevant documents.
+    Returns a formatted string of results.
+    """
+    results = COLL.query(
+        query_texts=[query],
+        n_results=max_results,
+        include=['documents', 'distances'],
     )
-    return agent
+    docs = results.get("documents", [])[0]
+    dists = results.get("distances", [])[0]
+    if not docs:
+        return "No relevant documents found."
+    output_lines = []
+    for i, (doc, dist) in enumerate(zip(docs, dists), 1):
+        output_lines.append(f"{i}. (score: {dist:.4f})\n{doc[:200]}...")
+    return "\n\n".join(output_lines)
 
-if __name__ == "__main__":
+# Create agent
-    ag = create_rag_agent()
+SYSTEM_PROMPT = (
-    # Simple demo loop
+    "You are an assistant that can search and add to a knowledge base. Use the provided tools."
-    while True:
+)
-        user_input = input("User: ")
+AGENT = create_agent(
-        if user_input.lower() in ("quit", "exit"):
+    llm=None,  # No LLM needed for tool calls; agent will use system prompt only
-            break
+    tools=[add_to_knowledge_base, search_knowledge_base],
-        result = ag.ainvoke(
+    system_prompt=SYSTEM_PROMPT,
-            {"messages": [HumanMessage(content=user_input)]},
+)
-            {"configurable": {"thread_id": "demo"}},
+
-        )
+# Expose a simple invoke function
-        print(result["messages"][-1].content)
+async def run_agent(messages: List[HumanMessage]):
+    return await AGENT.ainvoke({"messages": messages}, {"configurable": {"thread_id": "rag-agent"}})