Update rag_tools.py

rag_tools.py

@@ -1,80 +1,54 @@
-"""Tools for the RAG agent.
+"""
-
+Tools for the RAG agent: local semantic search and web search via Tavily.
-This module defines two LangChain tools:
-
-* ``search_local_kb`` – semantic search in the local ChromaDB vector store.
-* ``web_search`` – real‑time web search using Tavily.
-
-Both tools return a string containing the retrieved information.
 """
 
 from typing import List
 
-from langchain_ollama import ChatOllama
-from langchain_tavily import TavilySearchResults
 from langchain.tools import tool
+from langchain_ollama import ChatOllama
+from langchain_chroma import Chroma
+from langchain_tavily import TavilySearchResults
 
-# The LLM used for summarising or formatting responses
+# Local semantic search tool
-llm = ChatOllama(model="llama3")
-
-# Tavily client – the API key is read from the environment by the package
-# (requires a .env file or the TAVILY_API_KEY environment variable).
-search = TavilySearchResults()
-
-# ---------------------------------------------------------------------------
-# Local knowledge base search tool
-# ---------------------------------------------------------------------------
 @tool("search_local_kb")
-
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Perform a semantic search in the local ChromaDB vector store.
+    """Search the local ChromaDB knowledge base.
 
     Parameters
     ----------
     query: str
-        The user question.
+        The user's query.
     top_k: int, optional
-        Number of top documents to return. Defaults to 3.
+        Number of top results to return.
 
     Returns
     -------
     str
-        A formatted string containing the retrieved passages.
+        Concatenated content of the top results.
     """
-    # The vectorstore is expected to be loaded globally – the agent will
+    # Load the vector store (persisted)
-    # provide it via the tool context.  We simply call the retriever.
+    vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=ChatOllama(model="nomic-embed-text"))
-    retriever = globals().get("vectorstore_retriever")
+    retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    if retriever is None:
+    docs = retriever.invoke(query)
-        raise RuntimeError("Vector store retriever not configured for the tool.")
+    # docs is a list of Document objects
+    return "\n\n---\n\n".join([doc.page_content for doc in docs])
 
-    docs = retriever.get_relevant_documents(query, k=top_k)
+# Web search tool via Tavily
-    # Concatenate the documents into a single string.
-    passages = "\n\n".join(doc.page_content for doc in docs)
-    return passages
-
-# ---------------------------------------------------------------------------
-# Web search tool
-# ---------------------------------------------------------------------------
 @tool("web_search")
-
 def web_search(query: str) -> str:
-    """Search the web using Tavily and return the top results.
+    """Perform a web search using Tavily.
 
     Parameters
     ----------
     query: str
-        The user question.
+        The user's query.
 
     Returns
     -------
     str
-        A formatted string containing the search results.
+        Summarized search results.
     """
-    results = search.run(query)
+    tavily = TavilySearchResults(api_key="${TAVILY_API_KEY}")
-    # TavilySearchResults returns a list of dicts with keys: title, url, content
+    results = tavily.run(query)
-    formatted = []
+    # results is a list of dicts with keys: title, url, content
-    for r in results:
+    return "\n\n---\n\n".join([f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results])
-        formatted.append(f"Title: {r.get('title', 'N/A')}\nURL: {r.get('url', 'N/A')}\nSnippet: {r.get('content', 'N/A')}\n")
-    return "\n\n".join(formatted)
-
-# End of rag_tools.py