agent.py

########## Imports ############
from langchain_google_genai import ChatGoogleGenerativeAI
import os
from typing import TypedDict, List, Dict, Any, Optional
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage

from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
import string

from langchain_experimental.tools import PythonREPLTool
import ast, json

from langchain_community.tools import DuckDuckGoSearchRun
import os
from langchain_huggingface import HuggingFaceEndpoint , ChatHuggingFace
from langchain import LLMChain, PromptTemplate

def get_gpt_and_answer(question):
      # إنشاء LLM باستخدام endpoint من Hugging Face
      llm = HuggingFaceEndpoint(
          repo_id="openai/gpt-oss-20b",               # أو أي نموذج موجود في HF
          task="text-generation",

          huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
          provider="auto"              # يسمح باختيار المزود تلقائياً

      )

      # إنشاء سلسلة (chain) بسيطة
      template = """
      Question: {question}
      Answer:
      """
      llm = ChatHuggingFace(llm=llm)

      prompt = PromptTemplate.from_template(template)
      chain = LLMChain(llm=llm, prompt=prompt)

      # تجربة
      # result = chain.invoke({"question": question})

      model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
      result = (model.invoke([HumanMessage(content=template)]).content)

      return result #(result['text'])


########## State ############
class InfoState(TypedDict):
    question: str
    answer_type: Optional[str] # WebInfo  - WIKI - MATH
    answer_code : Optional[str]
    main_parts: Optional[List[str]]
    tool_answer : Optional[list[str]]
    final_answer : Optional[str]



######### Nodes ############
def get_wiki_relate(state: InfoState) -> InfoState:
    """
    Tool to Get the wikipedia info from keywords extracted from preprocessing at main_parts.

    Uses: Wikipedia API
    Returns: tool_answer (summary)
    """
    print("Using Wikipedia...")
    # Create the Wikipedia utility
    wiki = WikipediaAPIWrapper(
        lang="en",        # Wikipedia language
        top_k_results=1,  # how many results to fetch
        doc_content_chars_max=2000
    )

    # Make a tool from it
    wiki_tool = WikipediaQueryRun(api_wrapper=wiki)

    try:
      wiki_answer = wiki_tool.run(" ".join(state["main_parts"]) + " full wikipedia article about this topic")
      state['tool_answer'] = wiki_answer
      return state
    except Exception as e:
      print("Rate limit Exception")
      state['tool_answer'] = ""
      return state


def execute_code(state: InfoState) -> InfoState :
    """Tool to calculate any math using python code or get current date time."""
    print("Execut Code...")
    python_tool = PythonREPLTool()
    code = state["answer_code"]
    state["tool_answer"]=python_tool.run(code)
    return state

def get_code(state:InfoState) -> InfoState:
    """From prompt get the code to run."""
    print("Getting Code (Gemini)...")
    prompt = (
    f"You are a strict code generator. "
    f"Given the question: '{state['question']}', "
    f"return ONLY valid Python code that computes the answer IF the question is about math, date, or time. "
    f"Otherwise, return exactly: print('not valid')\n\n"
    f"Rules:\n"
    f"- Output ONLY the code or print('not valid')\n"
    f"- No explanations, no markdown, no extra text\n"
    f"- No quotes around the code\n"
    f"- Use print() to show the result\n"
    f"- Import modules only if needed (e.g. datetime, math)"
)

 
    state["answer_code"] = get_gpt_and_answer(prompt).strip()


    return state



def preprocess_text(state: dict) -> InfoState:

    """
    Preprocess text to get the keywords to help get results directly from wikipedia.

    Input: raw question
    Output: main_parts (list of keywords)
    """
    print("Preprocess text (Gemini)...")
    # 1️⃣ Prepare the prompt
    prompt =  (
    "We want to find the best-matching English Wikipedia pages for a factual question, "
    "so we must extract only the essential topic names or entities that Wikipedia likely has pages for. "
    "These should include the main subject (e.g., a person, event, place, or concept) and any directly relevant subtopic "
    "(like 'Discography', 'Career', or 'History') if they help narrow the search.\n\n"

    "Rules:\n"
    "- Output 1 to 3 items maximum.\n"
    "- Use English Wikipedia title format (capitalize each main word).\n"
    "- Translate non-English names or terms to English.\n"
    "- Exclude question words, pronouns, and filler terms.\n"
    "- Fix spelling errors if necessary.\n"
    "- Prefer specific Wikipedia topics over vague ones.\n\n"
    "Do NOT include markdown formatting or code blocks like ```json```. Output plain JSON only.\n"


    "Example:\n"
    "Q: 'Who built the Eiffel Tower?'\n"
    "A: [\"Eiffel Tower\", \"Gustave Eiffel\"]\n\n"
    f"Question: '{state['question']}'\n\n"
    "Output ONLY a valid JSON list as described — no explanations, markdown, or extra formatting."
    )


    response = get_gpt_and_answer(prompt).strip()

    # 3️⃣ Try to safely parse
    try:
        # First, try JSON
        state["main_parts"] = json.loads(response)
    except json.JSONDecodeError:
        try:
            # If not JSON, try Python literal
            state["main_parts"] = ast.literal_eval(response)
        except Exception:
            # If both fail, store fallback info
            print("⚠️ Model returned invalid content:", response)
            state["main_parts"] = []

    return state




def get_answer(state: InfoState) -> InfoState :
    """
    Final Node that returns the final answer organized.

    Combines: tool_answer → final_answer
    """
    print("Getting Answer (Gemini)...")

    prompt = (
        "You are a knowledgeable assistant that answers questions based on context and common factual knowledge.\n"
        "Use the context first, but if it clearly lacks the needed details, you may rely on well-known public facts "
        "(such as from Wikipedia) that logically complete the answer.\n\n"
        f"Question: {state['question']}\n"
        f"Context:\n{state.get('tool_answer')}\n\n"
        "Instructions:\n"
        "- Focus on producing one short factual answer.\n"
        "- You Should think enough at first before giving final answer. \n"
        "- Do not include tool names, prefixes, or metadata.\n"
        "- If the context contains partial hints, you can infer the answer from general knowledge of the same topic.\n"
        "- If the question asks about an attached file or audio, reply briefly that you cannot access attachments or audio files."
        "- If the context lacks key details or references a file, start with: I'm not sure because the question depends on missing data or an attached file. Then, add what you reasonably know about the topic."
        "- Final answer should be complete text not part of answer"
        "Final Answer:"
    )

    state["final_answer"] = get_gpt_and_answer(prompt)

    return state
    
def get_type(state: InfoState) -> InfoState:
    """Choose which tool to use based on question type (WIKI, SEARCH, CODE)."""

    print("Getting Type (Gemini/OpenAI)...")
    q = state["question"].lower()
    if any(w in q for w in ["calculate", "sum", "add", "price", "how much", "date", "time"]):
        state["answer_type"] = "MATH"
    elif any(w in q for w in ["latest", "current", "today", "news"]):
        state["answer_type"] = "WebInfo"
    elif any(w in q for w in ["who", "where", "what", "when", "why", "how many", "which"]):
        state["answer_type"] = "WIKI"
    else:
        try:
            prompt = f"""
            You are a strict classifier.
            Your job is to classify the following question into ONE of four categories:
        
            - WIKI → informative, factual, or science question
            - WebInfo → up-to-date, news, or current event question
            - MATH → math, numeric, date, or time calculation
            - LLM → all others, including links, reasoning, and file-related tasks
        
            Question: "{state['question']}"
        
            Rules:
            - Reply with exactly one of these words: WIKI, WebInfo, MATH, or LLM
            - Output nothing else. No punctuation, no quotes, no explanation.
            - If unsure, default to LLM.
            """
            resp = get_gpt_and_answer(prompt)
            state["answer_type"] = resp.split()[0].strip()
        except:
            state["answer_type"] = "LLM"
    return state




def get_search_results(state: InfoState) -> InfoState:
    """Tool to search web for results using DuckDuckGo."""
    print("Searching...")

    search = DuckDuckGoSearchRun()

    try:
      state['tool_answer'] = search.run(state["question"])   #" " .join(state["main_parts"]))
      return state
    except Exception:
      state['tool_answer'] = ""
      return state


def route(state: InfoState):
    print(state["answer_type"])
    return state["answer_type"]



################# Graph ################
def get_graph():
        graph = StateGraph(InfoState)

        # Add nodes
        #graph.add_node("get_wiki_relate", get_wiki_relate)
        graph.add_node("preprocess_text", preprocess_text)
        graph.add_node("get_answer", get_answer)
        graph.add_node("get_type", get_type)
        graph.add_node("get_search_results", get_search_results)
        graph.add_node("execute_code", execute_code)
        graph.add_node("get_code", get_code)

        # Add edges
        graph.add_edge(START, "preprocess_text")
        graph.add_edge("preprocess_text", "get_type")


        # Add conditional edges
        graph.add_conditional_edges(
            "get_type",
            route,
            {
                "WebInfo": "get_search_results",
                "WIKI": "get_search_results",#"get_wiki_relate",
                "MATH": "get_code",
                "LLM": "get_answer"
            }
        )

        # Add final edges
        graph.add_edge("get_search_results", "get_answer")
        #graph.add_edge("get_wiki_relate", "get_answer")

        graph.add_edge("get_code", "execute_code")
        graph.add_edge("execute_code", "get_answer")

        graph.add_edge("get_answer", END)


        # Compile the graph
        compiled_graph = graph.compile()
        return compiled_graph

def ask(compiled_graph,question):
        legitimate_result = compiled_graph.invoke({
            "question": question,

        })

        return legitimate_result['final_answer'] #,legitimate_result