Update rag_tools.py

rag_tools.py

@@ -2,90 +2,62 @@
 
 Two tools are provided:
 
-* :func:`search_local_kb` – semantic search in the local ChromaDB vector store.
-* :func:`web_search` – real‑time web search via Tavily.
-
-Both functions are decorated with :func:`langchain.tools.tool` so that LangChain
-can expose them to the agent.
+* ``search_local_kb`` – semantic search in the local ChromaDB store.
+* ``web_search`` – real‑time web search via Tavily.
 """
 
-from __future__ import annotations
-
-import os
 from typing import List
 
-from langchain.chains import RetrievalQA
-from langchain.chroma import Chroma
 from langchain_ollama import ChatOllama
 from langchain_tavily import TavilySearchResults
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
+
 from langchain.tools import tool
 
-# Load environment variables (TAVILY_API_KEY)
-from dotenv import load_dotenv
+# The LLM used for generating answers.  Using the same model as the embeddings
+# keeps the pipeline consistent.
+_llm = ChatOllama(model="llama3")
+
+# Tavily client – the API key is read from the environment variable
+# ``TAVILY_API_KEY`` by the TavilySearchResults class.
+_tavily = TavilySearchResults()
 
-load_dotenv()
 
-# ---------------------------------------------------------------------------
-# Local KB search
-# ---------------------------------------------------------------------------
 @tool("search_local_kb")
-def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma | None = None) -> str:
-    """Return the best answer from the local ChromaDB store.
+def search_local_kb(query: str, top_k: int = 3, vectorstore: Chroma = None) -> str:
+    """Semantic search in the local ChromaDB vector store.
 
-    Parameters
-    ----------
-    query: str
-        The user question.
-    top_k: int
-        Number of documents to retrieve.
-    vectorstore: Chroma | None
-        If ``None`` the function will attempt to load the default store.
+    Args:
+        query: User question.
+        top_k: Number of results to return.
+        vectorstore: The Chroma instance to query.
 
-    Returns
-    -------
-    str
-        The answer prefixed with the source identifier.
+    Returns:
+        A string containing the concatenated top results.
     """
     if vectorstore is None:
-        # Load the default store
-        from vectorstore import create_vectorstore
-        vectorstore = create_vectorstore()
-
+        raise ValueError("vectorstore must be provided")
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=ChatOllama(model="llama3", temperature=0),
-        chain_type="stuff",
-        retriever=retriever,
-    )
-    result = qa_chain.run(query)
-    return f"[chromadb] {result}"
+    docs: List[Document] = retriever.invoke(query)
+    return "\n\n".join(doc.page_content for doc in docs)
+
 
-# ---------------------------------------------------------------------------
-# Web search via Tavily
-# ---------------------------------------------------------------------------
 @tool("web_search")
 def web_search(query: str, top_k: int = 3) -> str:
-    """Return an answer based on a Tavily web search.
+    """Search the web using Tavily.
 
-    Parameters
-    ----------
-    query: str
-        The user question.
-    top_k: int
-        Number of search results to retrieve.
+    Args:
+        query: User question.
+        top_k: Number of results to return.
 
-    Returns
-    -------
-    str
-        The answer prefixed with the source identifier.
+    Returns:
+        Concatenated snippets from the search results.
     """
-    api_key = os.getenv("TAVILY_API_KEY")
-    if not api_key:
-        raise RuntimeError("TAVILY_API_KEY environment variable is not set")
-    tavily = TavilySearchResults(api_key=api_key, top_k=top_k)
-    results = tavily.run(query)
-    snippets = "\n\n".join(item.get("content", "") for item in results)
-    llm = ChatOllama(model="llama3", temperature=0)
-    prompt = f"Answer the question based on the following web snippets:\n{snippets}\n\nQuestion: {query}\nAnswer:"  # noqa: E501
-    answer = llm.invoke(prompt).content
-    return f"[tavily] {answer}"
+    results = _tavily.run(query, max_results=top_k)
+    # TavilySearchResults returns a list of dicts with keys like 'title',
+    # 'content', 'url'.  We return the content for simplicity.
+    return "\n\n".join(r.get("content", "") for r in results)
+
+*** End of File ***