Source code for orkgnlp.deepresearch.config

# -*- coding: utf-8 -*-
from typing import List, Optional

import tiktoken
from openai import OpenAI




class ModelConfig:
    """Configuration class for OpenAI models with structured output support."""

    MIN_CHUNK_SIZE = 140

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.openai.com/v1",
        custom_model: Optional[str] = None,
    ):
        self.custom_model = custom_model
        self.openai_client = OpenAI(
            api_key=api_key,
            base_url=base_url,
        )

        # Initialize tokenizer
        try:
            self.encoder = tiktoken.get_encoding("o200k_base")
        except Exception:
            # Fallback to cl100k_base if o200k_base is not available
            self.encoder = tiktoken.get_encoding("cl100k_base")



    def get_model_config(self) -> dict:
        """Get the appropriate model configuration."""
        if not self.openai_client:
            raise Exception("No OpenAI client available - check OPENAI_KEY")

        model = self.custom_model

        if model:
            return {
                "client": self.openai_client,
                "model": model,
                "structured_outputs": True,
                "reasoning_effort": None,
            }

        # Default to o3-mini with reasoning effort
        return {
            "client": self.openai_client,
            "model": "o3-mini",
            "structured_outputs": True,
            "reasoning_effort": "medium",
        }




    def generate_completion(
        self, messages: List[dict], response_format: Optional[dict] = None, **kwargs
    ) -> dict:
        """Generate a completion using the configured model."""
        config = self.get_model_config()

        # Prepare parameters
        params = {"model": config["model"], "messages": messages, **kwargs}

        # Add structured output if provided
        if response_format and config["structured_outputs"]:
            params["response_format"] = response_format

        # Add reasoning effort for o3 models
        if config["reasoning_effort"] and "o3" in config["model"]:
            params["reasoning_effort"] = config["reasoning_effort"]

        return config["client"].chat.completions.create(**params)