src/internal/agents/api_tool_agent.py

from src.internal.agents.base_agents import Agent, AgentRequest
from src.internal.agents.gpt_executor_agent import GPTExecutorAgent
from src.domain.api_fetch_types import APIRequest
from src.utils.json_cleaner import clean_response
from typing import List, Dict
from openai import OpenAI
import json
from dataclasses import dataclass
import logging
import requests
class APIToolAgent(Agent):
    def __init__(self, executor_agent : Agent, api_host : str = "http://localhost"):
        super().__init__(None)
        self.executor_agent = executor_agent
        self.api_host = api_host

    async def api_fetch(self, req:AgentRequest):
        agent_result = ""
        api_response = None
        async for item in self.executor_agent.get_result(req):
            if item["type"] == "chunk":
                    chunk = item["data"]["chunk"]
                    logging.info(chunk)
                    agent_result += chunk

            elif item["type"] == "metadata":
                setup_time = item['data']['setup_time']
                print(f"\nSetup completed in {setup_time:.2f}s")
                    
            elif item["type"] == "complete":
                total_time = item['data']['total_time']
                print(f"\nTotal time: {total_time:.2f}s")

        logging.info(f"agent result {agent_result}")
        print(f"agent result {agent_result}")
        cleaned_agent_result = clean_response(agent_result)

        if cleaned_agent_result and ("error" not in agent_result):
            api_req = APIRequest(
                endpoint = cleaned_agent_result["endpoint"],
                method = cleaned_agent_result["method"],
                request = cleaned_agent_result["request"]
            )
            if(api_req.method == "POST"):
                api_response = requests.post(self.api_host + "/" + api_req.endpoint, json = api_req.request)
            elif(api_req.method == "GET"):
                api_response = requests.get(self.api_host + "/" + api_req.endpoint, json = api_req.request)
            api_response = str(api_response.json())
            yield str(api_response)
        else:
            yield str(agent_result)

    async def get_result(self, req:AgentRequest):
        logging.info(f"Question = {req.question}")
        print("Question = ", req.question)
        async for item in self.api_fetch(req):
            yield item