api/researcher.py

# Handles API logic for calling scraper & LLM
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, HttpUrl
from llm_module.message import chain
from scraper import multiple_links_scraping
from typing import Any
import logging

logger = logging.getLogger('fastapi_cli')
api_router = APIRouter(prefix='/researcher', tags=['Researcher'])

class InputField(BaseModel):
    url : str
    generic_message : str

class OutputField(BaseModel):
    company_summary : str
    customised_message : str


@api_router.post('/invoke', response_model=OutputField)
async def invoke_web_researcher(body: InputField):
    try:
        scraped_data = await multiple_links_scraping([body.url])
    except Exception as e:
        logger.error(f"Failed to scrape: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")
    try:
        output = await chain.ainvoke({
            "generic_message": body.generic_message,
            "scraped_message": scraped_data[0]
        })

        if output is None :
            raise ValueError("LLM chain returned None values")
        print(output)
        return OutputField(
            company_summary=output['summary_output'],
            customised_message=output['generated_message']
        )
    except Exception as e:
        logger.error(f"Failed to invoke the chain: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")