|
|
import os |
|
|
from langchain_core.tools import tool |
|
|
from pydantic import BaseModel, Field |
|
|
from rag_with_memory import MemoryRAG |
|
|
from loguru import logger |
|
|
|
|
|
|
|
|
_rag_engine = None |
|
|
|
|
|
def get_rag_engine(): |
|
|
global _rag_engine |
|
|
if _rag_engine is None: |
|
|
|
|
|
target_path = "/app/data/knowledge_base" |
|
|
|
|
|
|
|
|
if not os.path.exists(target_path) or not os.listdir(target_path): |
|
|
logger.error(f"KNOWLEDGE BASE MISSING at {target_path}") |
|
|
return None |
|
|
|
|
|
try: |
|
|
logger.info("Initializing MemoryRAG engine...") |
|
|
_rag_engine = MemoryRAG(docs_path=target_path) |
|
|
logger.success("MemoryRAG engine is now ONLINE.") |
|
|
except Exception as e: |
|
|
logger.error(f"MemoryRAG Initialization failed: {e}") |
|
|
return None |
|
|
return _rag_engine |
|
|
|
|
|
class KnowledgeBaseInput(BaseModel): |
|
|
query: str = Field(description="The search query.") |
|
|
|
|
|
@tool(args_schema=KnowledgeBaseInput, return_direct=True) |
|
|
def knowledge_base_search(query: str) -> str: |
|
|
"""Search the company knowledge base for answers.""" |
|
|
engine = get_rag_engine() |
|
|
|
|
|
if engine is None: |
|
|
logger.error("Tool called but RAG engine remains None.") |
|
|
return "I'm sorry, I'm currently unable to access my internal documentation." |
|
|
|
|
|
try: |
|
|
|
|
|
result = engine.query(query, session_id="tool_session") |
|
|
return result.get("answer", "No information found.") |
|
|
except Exception as e: |
|
|
logger.error(f"Tool execution failed: {e}") |
|
|
return "I encountered an error while searching." |
|
|
|