agent.py

import os
from dotenv import load_dotenv
from langchain.docstore.document import Document
from langchain_community.retrievers import BM25Retriever
from langchain.tools import Tool
from langchain.utilities import SerpAPIWrapper
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_core.messages import AnyMessage, HumanMessage
from langchain_groq import ChatGroq
from typing import TypedDict, Annotated
import fitz  # PyMuPDF

# Load environment variables
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")
serpapi_api_key = os.getenv("SERPAPI_API_KEY")


# --- PDF parsing ---
def parse_pdfs(uploaded_files):
    pdf_docs = []
    for uploaded_file in uploaded_files:
        with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:
            text = ""
            for page in doc:
                text += page.get_text()
            pdf_docs.append(Document(page_content=text, metadata={"source": uploaded_file.name}))
    return pdf_docs


# --- BM25 Retrieval ---
def build_retriever(all_docs):
    return BM25Retriever.from_documents(all_docs)


def extract_text(query: str, retriever):
    results = retriever.invoke(query)
    if results:
        return "\n\n".join([doc.page_content for doc in results[:3]])
    else:
        return "لم يتم العثور على معلومات مطابقة في الملفات."


# --- Additional Tools --- 

def calculator_tool_func(query: str):
    try:
        # أداة حساب بسيطة باستخدام eval (يمكن تطويرها لاحقًا)
        result = str(eval(query, {"__builtins__": {}}))
        return result
    except Exception:
        return "تعذر حساب التعبير المدخل."

calculator_tool = Tool(
    name="Calculator",
    func=calculator_tool_func,
    description="Performs simple arithmetic calculations."
)

def weather_tool_func(location: str):
    # يمكن ربطها ب API حقيقي للطقس لاحقًا
    return f"حالة الطقس في {location}: مشمس، درجة الحرارة 25 درجة مئوية."

weather_tool = Tool(
    name="Weather",
    func=weather_tool_func,
    description="Provides weather information for a given location."
)


# --- Create NINU Agent ---
def create_ninu_agent(user_docs=None):
    bm25_retriever = build_retriever(user_docs) if user_docs else None

    def pdf_tool_func(q):
        if bm25_retriever:
            return extract_text(q, bm25_retriever)
        else:
            return "لا توجد ملفات PDF مرفوعة للبحث."

    NINU_tool = Tool(
        name="NINU_Lec_retriever",
        func=pdf_tool_func,
        description="Retrieves content from uploaded PDFs based on a query."
    )

    serpapi = SerpAPIWrapper(serpapi_api_key=serpapi_api_key)
    SerpAPI_tool = Tool(
        name="WebSearch",
        func=serpapi.run,
        description="Searches the web for recent information."
    )

    # دمج جميع التولز
    tools = [NINU_tool, SerpAPI_tool, calculator_tool, weather_tool]

    llm = ChatGroq(model="deepseek-r1-distill-llama-70b", groq_api_key=groq_api_key)
    llm_with_tools = llm.bind_tools(tools)

    class AgentState(TypedDict):
        messages: Annotated[list[AnyMessage], add_messages]

    def assistant(state: AgentState):
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

    builder = StateGraph(AgentState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")
    return builder.compile()


# --- Main interaction function ---
def run_ninu(query, user_docs=None):
    agent = create_ninu_agent(user_docs)

    conversation = []

    intro_prompt = """
You are a general AI assistant with access to several tools:
1. NINU_Lec_retriever: retrieves content from uploaded PDFs based on a query.
2. WebSearch: performs web searches to answer questions about current events or general knowledge.
3. Calculator: performs arithmetic calculations.
4. Weather: provides weather information for given locations.
Based on the user's query, decide whether to use one or more of these tools.
When answering, report your thoughts and finish your answer with the following template:
FINAL ANSWER: [YOUR FINAL ANSWER].
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
If you are asked for a number, don't use commas or units (like $, %, etc.) unless specified.
If you are asked for a string, avoid articles, abbreviations, and write digits in plain text unless specified.
"""
    conversation.append(HumanMessage(content=intro_prompt))
    conversation.append(HumanMessage(content=query))

    response = agent.invoke({"messages": conversation})
    return response["messages"][-1].content


# تسجيل الوكيل باسم `ninu` ليتم استيراده من ملفات أخرى
ninu = create_ninu_agent