File size: 3,020 Bytes
8a27bb1 8f6bb93 8a27bb1 ee3f4ca 8a27bb1 41e79c8 8a27bb1 41e79c8 8a27bb1 8f6bb93 ee3f4ca 41e79c8 ee3f4ca 41e79c8 8a27bb1 8f6bb93 2b63102 98714ab 0b87551 2b63102 0b87551 8a27bb1 41e79c8 8a27bb1 ee3f4ca 8a27bb1 41e79c8 ee3f4ca 8f6bb93 ee3f4ca 8a27bb1 8f6bb93 0b87551 2b63102 8f6bb93 0b87551 8f6bb93 0b87551 2b63102 0b87551 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import os
from functools import lru_cache
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.tools import tool
from langchain_community.tools import DuckDuckGoSearchRun
from langgraph.prebuilt import create_react_agent
from dotenv import load_dotenv
from src.rag_engine import KnowledgeBase
from src.file_processor import FileProcessor
load_dotenv()
file_processor = FileProcessor()
_fallback_kb = KnowledgeBase(pdf_path=os.path.join("data", "policy.pdf"))
try:
_fallback_kb.load_and_index()
except Exception as e:
print(f"Fallback KB skipped: {e}")
_search_tool = DuckDuckGoSearchRun()
# gemini-2.5-flash reliably answers after tool calls; override per key via env.
MODEL_NAME = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
SYSTEM_PROMPT = (
"You are a precise research assistant with two tools:\n"
"- `lookup_documents`: searches the user's uploaded files and the internal knowledge base. "
"Prefer this for questions about policies, documents, or any uploaded content.\n"
"- `search_web`: searches the live web. Use this for current events, news, or general "
"knowledge that is unlikely to be in the documents.\n\n"
"Guidelines:\n"
"1. Choose the tool that best fits the question; use both if needed.\n"
"2. Ground your answer in the retrieved content and do not invent facts. If the documents "
"do not contain the answer, say so and try the web.\n"
"3. Use the conversation summary and recent turns to resolve follow-up questions "
"(e.g. pronouns like 'it' or 'that').\n"
"4. Be concise, accurate, and cite the source of your information when relevant."
)
@tool
def lookup_documents(query: str) -> str:
"""Search the user-uploaded documents for relevant information.
Use this for questions about content in any uploaded files."""
if file_processor.has_documents():
result = file_processor.retrieve(query)
if result:
return result
return _fallback_kb.retrieve(query)
@tool
def search_web(query: str) -> str:
"""Search the web for current events, news, or general knowledge not in uploaded documents."""
try:
return _search_tool.run(query)
except Exception as e:
return f"Search failed: {e}"
@lru_cache(maxsize=32)
def get_llm(api_key: str) -> ChatGoogleGenerativeAI:
"""Cached Gemini client per key, reused by the agent and the summarizer."""
if not api_key or not api_key.strip():
raise ValueError("A Google Gemini API key is required.")
return ChatGoogleGenerativeAI(
model=MODEL_NAME,
temperature=0,
google_api_key=api_key.strip(),
)
@lru_cache(maxsize=32)
def get_agent_executor(api_key: str):
"""ReAct agent for a given key (BYOK). Tools and the index are shared; only the
LLM is per-key, and the graph is cached so repeat calls don't rebuild it."""
return create_react_agent(
get_llm(api_key),
[lookup_documents, search_web],
prompt=SYSTEM_PROMPT,
)
|