api/routers/orchestration.py

import logging
from fastapi import APIRouter
from api.stored_data import stored_data
from src.genai.orchestration_agent.agent import OrchestrationAgent
# from src.genai.orchestration_agent.utils.utils import handle_tools
from api.schemas.orchestration import OrchestrationRequest
from src.genai.orchestration_agent.utils.utils import handle_tools

app_logger = logging.getLogger("app_logger")
error_logger = logging.getLogger("error_logger")

router= APIRouter()
agent = OrchestrationAgent()
graph = agent.orchestration_graph()

@router.post("/orchestration", description="Gives the analytics of influencers.")
def orchestration_endpoint(request:OrchestrationRequest):
    try:
        config={"configurable": {"thread_id": "orchestration-thread"}}
        result = graph.invoke({'messages':[request.message],
                               'image_base64': request.image_base64}, config=config)
        print('Result tools:', result['tools'])
        tools=handle_tools(result['tools'], stored_data)
        print('Tools:',tools)
        if result['image_caption'] is not None: stored_data['image_caption']=result['image_caption']
        # if result['video_idea'] !='' or result.video_idea != 'null': stored_data['refined_ideation']= result['video_idea']
        # if result['video_story']!='' or result.video_story!='null': stored_data['final_story']= result['video_story']
        app_logger.info('Orchestrator executed')
            
        return {'tool_response':tools,
                'message_response': result['query_response'],}

    except Exception as e:
        error_logger.error(f'Unable to run orchestration: {e}')
        raise