utils.py

import numpy as np
import requests
import json


class HuggingFaceClient:
    def __init__(self, hf_token: str, model: str):
        self.base_url = "https://router.huggingface.co/v1"
        self.hf_token = hf_token
        self.model = model

    def generate_embedding(self, texts: str) -> np.ndarray:
        api_url = f"https://router.huggingface.co/hf-inference/models/{self.model}/pipeline/feature-extraction"
        headers = {"Authorization": f"Bearer {self.hf_token}"}
        response = requests.post(
            api_url,
            headers=headers,
            json={"inputs": texts, "options": {"wait_for_model": True}},
        )
        if response.status_code == 200:
            return response.json()
        """
        ollama_url = "http://localhost:11434/api/embed"
        headers = {"Content-Type": "application/json"}
        data = {"model": "embeddinggemma", "input": text}
        response = requests.post(
            ollama_url,
            data=json.dumps(data),
            headers=headers,
        )
        if response.status_code == 200:
            embedding = response.json().get("embeddings", [])
            return embedding[0]
        else:
            print(
                "Failed to generate embedding:",
                response.status_code,
                response.text,
            )
            return np.zeros(768)  # Return a zero vector on failure
        """


class OpenRouterClient:
    def __init__(self, model: str, api_key: str):
        self.model = model
        self.api_key = api_key
        self.base_url = "https://openrouter.ai/api/v1"

    def chat(self, prompt: str) -> str:
        response = requests.post(
            url="https://openrouter.ai/api/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {self.api_key}",
            },
            data=json.dumps(
                {
                    "model": self.model,
                    "messages": [
                        {
                            "role": "user",
                            "content": prompt,
                        }
                    ],
                }
            ),
        )
        response.raise_for_status()
        j = response.json()["choices"][0]["message"]["content"]
        return j