bagoodex_client.py

import os
import requests
from openai import OpenAI
from dotenv import load_dotenv
from r_types import ChatMessage
from prompts import SYSTEM_PROMPT_BASE, SYSTEM_PROMPT_MAP
from typing import List

load_dotenv()
API_KEY = os.getenv("AIML_API_KEY")
API_URL = "https://api.aimlapi.com"

class BagoodexClient:
    def __init__(self, api_key=API_KEY, api_url=API_URL):
        self.api_key = api_key
        self.api_url = api_url
        self.client = OpenAI(base_url=self.api_url, api_key=self.api_key)

    def complete_chat(self, query):
        """
        Calls the standard chat completion endpoint using the provided query.
        Returns the generated followup ID and the text response.
        """
        response = self.client.chat.completions.create(
            model="bagoodex/bagoodex-search-v1",
            messages=[
                ChatMessage(role="user", content=SYSTEM_PROMPT_BASE), 
                ChatMessage(role="user", content=query)
            ],
        )
        followup_id = response.id  # the unique ID for follow-up searches
        answer = response.choices[0].message.content
        return followup_id, answer
    
    def base_qna(self, messages: List[ChatMessage], system_prompt=SYSTEM_PROMPT_BASE):
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                ChatMessage(role="user", content=system_prompt), 
                *messages
            ],
        )
        return response.choices[0].message.content

    def get_links(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(
            f"{self.api_url}/v1/bagoodex/links", headers=headers, params=params
        )
        return response.json()

    def get_images(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(
            f"{self.api_url}/v1/bagoodex/images", headers=headers, params=params
        )
        return response.json()

    def get_videos(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(
            f"{self.api_url}/v1/bagoodex/videos", headers=headers, params=params
        )
        return response.json()

    def get_local_map(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(
            f"{self.api_url}/v1/bagoodex/local-map", headers=headers, params=params
        )
        return response.json()

    def get_knowledge(self, followup_id):
        headers = {"Authorization": f"Bearer {self.api_key}"}
        params = {"followup_id": followup_id}
        response = requests.get(
            f"{self.api_url}/v1/bagoodex/knowledge", headers=headers, params=params
        )
        return response.json()