| import os
|
| from langchain_community.vectorstores import Chroma
|
| from langchain_community.embeddings import HuggingFaceEmbeddings
|
| from langchain_core.tools import tool
|
|
|
| CHROMA_DIR = "chroma_store"
|
|
|
| _store = None
|
|
|
| def _get_store():
|
| global _store
|
| if _store is None:
|
| embeddings = HuggingFaceEmbeddings(
|
| model_name="sentence-transformers/all-MiniLM-L6-v2",
|
| model_kwargs={"device": "cpu"},
|
| encode_kwargs={"normalize_embeddings": True},
|
| )
|
|
|
| _store = Chroma(
|
| persist_directory=CHROMA_DIR,
|
| embedding_function=embeddings,
|
| collection_name="research_docs",
|
| )
|
| return _store
|
|
|
| @tool
|
| def retriever_tool(query: str) -> dict:
|
| """Search the indexed documents for information relevant to the query. Always try this first before web search."""
|
| try:
|
| store = _get_store()
|
| docs = store.similarity_search(query, k = 5)
|
| if not docs:
|
| return {"found": False, "message": "Nothing found in documents."}
|
| return {
|
| "found": True,
|
| "results": [
|
| {
|
| "content" : d.page_content,
|
| "source": d.metadata.get("source_file", "unknown"),
|
| }
|
| for d in docs
|
| ],
|
| }
|
| except Exception as e:
|
| return {"found": False, "error": str(e)}
|
|
|
|
|
| def get_retriever():
|
| return _get_store().as_retriever(search_kwargs={"k": 5}) |
