my_agent.py

import os

from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph import StateGraph, START, MessagesState
from langchain.agents import create_agent
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_ollama import ChatOllama
from langchain.agents.middleware.types import AgentState
from langchain.messages import HumanMessage, AIMessage, SystemMessage
from langfuse.langchain import CallbackHandler

import add_telemetry  # noqa: F401 to initialize Langfuse telemetry

hf_token = os.getenv("HF_TOKEN")
langfuse_handler = CallbackHandler()


class AgentResponseState(AgentState):
    response: str
    messages: list[HumanMessage | AIMessage]


# --- Basic Agent Definition ---
# ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
class BasicAgent:
    def __init__(self):
        model = HuggingFaceEndpoint(
            repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
            task="text-generation",
            max_new_tokens=512,
            do_sample=False,
            repetition_penalty=1.03,
        )
        llm = ChatHuggingFace(llm=model, verbose=True)
        # llm = ChatOllama(
        #     model="qwen3:0.6b",
        #     api_base="http://localhost:11434",  # replace with
        #     # debug=True,
        # )
        tools = [
            DuckDuckGoSearchRun(),
        ]

        builder = StateGraph(MessagesState)

        model = create_agent(llm, tools)
        builder.add_node("assistant", model)
        builder.add_node("tools", ToolNode(tools))

        # Define edges: these determine how the control flow moves
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges(
            "assistant",
            # If the latest message requires a tool, route to tools
            # Otherwise, provide a direct response
            tools_condition,
        )
        builder.add_edge("tools", "assistant")
        self.agent = builder.compile()

        print("BasicAgent initialized.")

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")

        fixed_answer = self.generate_answer(question)

        print(f"Agent returning fixed answer: {fixed_answer}")
        return fixed_answer

    def generate_answer(self, question: str) -> str:
        response = self.agent.invoke(
            {
                "messages": [
                    {
                        "role": "user",
                        "content": question,
                    }
                ]
            }
        )
        print(f"Agent raw response: {response}")
        print(f"response.content => {response['messages']}")
        print(f"AI response => {response['messages'][-1].content}")
        return response['messages'][-1].content