1+ from langchain_openai import ChatOpenAI
2+ from abi .services .agent .Agent import Agent , AgentConfiguration , AgentSharedState
3+ from src import secret
4+ from fastapi import APIRouter
5+ from src .marketplace .__demo__ .workflows .ExecutePythonCodeWorkflow import ExecutePythonCodeWorkflow , ExecutePythonCodeWorkflowConfiguration
6+ from enum import Enum
7+ from typing import Optional
8+ from pydantic import SecretStr
9+ from langchain_core .tools import Tool , BaseTool
10+
11+ NAME = "BodoAgent"
12+ MODEL = "gpt-4o"
13+ TEMPERATURE = 0
14+ DESCRIPTION = "An agent that can analyze large data with Bodo DataFrames"
15+ # TODO: Add avatar
16+ AVATAR_URL = "https://raw.githubusercontent.com/jupyter-naas/awesome-notebooks/refs/heads/master/.github/assets/logos/Naas.png"
17+ SYSTEM_PROMPT = f"""
18+ You are { NAME }
19+ You can execute Python code through the ExecutePythonWorkflow tool to perform your analyses.
20+
21+ When a user asks a question involving data (e.g., describing a dataset, computing aggregates, or exploring patterns), you should:
22+
23+ * Write a complete Python script to perform the analysis.
24+
25+ * Always import bodo.pandas as pd at the top of the script (never use regular pandas).
26+
27+ * Read data using pandas-style APIs (for example: `pd.read_parquet("/path/to/file"))`.
28+
29+ * Perform any operations or computations using standard pandas syntax (groupby, describe, value_counts, etc.).
30+
31+ * Print concise, readable outputs — summary statistics, shapes, missing value counts, or aggregates.
32+
33+ * Avoid unsafe or network operations; only read data and compute results.
34+
35+ * After running the workflow, summarize the findings in plain English.
36+
37+ Your responses should be short, factual, and focused on analytical insights rather than speculation.
38+ """
39+ SUGGESTIONS = ["Summarize this CSV file: /path/to/file.csv" ]
40+
41+ def create_agent (
42+ agent_shared_state : Optional [AgentSharedState ] = None ,
43+ agent_configuration : Optional [AgentConfiguration ] = None
44+ ) -> Agent :
45+ # Init
46+ tools : list [Tool | BaseTool | Agent ] = []
47+
48+ # Set model
49+ model = ChatOpenAI (
50+ model = MODEL ,
51+ temperature = TEMPERATURE ,
52+ api_key = SecretStr (secret .get ('OPENAI_API_KEY' ))
53+ )
54+
55+ # Set configuration
56+ if agent_configuration is None :
57+ agent_configuration = AgentConfiguration (system_prompt = SYSTEM_PROMPT )
58+ if agent_shared_state is None :
59+ agent_shared_state = AgentSharedState (thread_id = "0" )
60+
61+ # Add tools
62+ config = ExecutePythonCodeWorkflowConfiguration (timeout = 600 , allow_imports = True )
63+ tools += ExecutePythonCodeWorkflow (config ).as_tools ()
64+
65+ return BodoAgent (
66+ name = NAME ,
67+ description = DESCRIPTION ,
68+ chat_model = model ,
69+ tools = tools ,
70+ agents = [],
71+ state = agent_shared_state ,
72+ configuration = agent_configuration ,
73+ # memory is automatically configured based on POSTGRES_URL environment variable
74+ )
75+
76+ class BodoAgent (Agent ):
77+ def as_api (
78+ self ,
79+ router : APIRouter ,
80+ route_name : str = NAME ,
81+ name : str = NAME ,
82+ description : str = "API endpoints to call the Bodo agent completion." ,
83+ description_stream : str = "API endpoints to call the Bodo agent stream completion." ,
84+ tags : Optional [list [str | Enum ]] = None ,
85+ ):
86+ return super ().as_api (router , route_name , name , description , description_stream , tags )
0 commit comments