File size: 2,773 Bytes
4533b99 ac32126 4533b99 2fce62e ac32126 4533b99 2fce62e 4533b99 2fce62e 4533b99 2fce62e ac32126 2fce62e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import os
from typing import TypedDict, List, Dict, Any, Optional
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, AIMessage
from langgraph.graph.message import add_messages
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
import sys
from pathlib import Path
from pydantic import BaseModel
import json
class OneResponse(BaseModel):
question_id: str
answer: str
root_path = Path(__file__).parent.parent.parent
print(f"Root path for imports: {root_path}")
sys.path.insert(0, str(Path(__file__).parent.parent))
from agent.tools.search_tool import search_tool
# --- zBottaAgent Definition ---
class zBottaAgent:
def __init__(self):
print("zBottaAgent initialized.")
chat = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0.7)
tools = [search_tool]
response_schema = OneResponse.model_json_schema()
self.chat_with_tools = chat.bind(
tools=tools,
response_mime_type="application/json",
response_schema=response_schema,
)
def __call__(self, question: str) -> str:
print(f"Agent received question (first 50 chars): {question[:50]}...")
messages = [
SystemMessage(content="You are a general AI assistant. I will ask you a question. " \
"Report your thoughts, and finish your answer. " \
"Your 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 comma to write your number neither use " \
"units such as $ or percent sign unless specified otherwise. If you are asked for a string, " \
"don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text " \
"unless specified otherwise. If you are asked for a comma separated list, apply the above rules " \
"depending of whether the element to be put in the list is a number or a string."),
HumanMessage(content=question)
]
response = self.chat_with_tools.invoke(messages)
print(f"Agent raw response: {response}")
answer = json.loads(response.content)
final_answer = answer["answer"]
print(f"Final answer: {final_answer}")
return final_answer
if __name__ == "__main__":
os.environ["GOOGLE_API_KEY"] = "AIzaSyA28oI_EEL3Io4I1OGF_Yj890ioKcCKWlo" # Set your Google API key here for testing
agent = zBottaAgent()
test_question = "What is the capital of France?"
print(f"Testing agent with question: {test_question}")
answer = agent(test_question) |