Mohammad Wasil
final update
e1550b3
import os
from langchain_core.tools import tool
from pydantic import BaseModel, Field
from rag_with_memory import MemoryRAG
from loguru import logger
# Global variable to hold the engine
_rag_engine = None
def get_rag_engine():
global _rag_engine
if _rag_engine is None:
# ABSOLUTE PATH for Docker 2026
target_path = "/app/data/knowledge_base"
# Verify folder exists and isn't empty
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:
# LOGIC FIX: Always return a string for the Agent to read
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."