diff --git a/kura/cluster.py b/kura/cluster.py index 7e5b1e3..5227142 100644 --- a/kura/cluster.py +++ b/kura/cluster.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from instructor.models import KnownModelName from instructor import AsyncInstructor + import instructor import numpy as np import asyncio @@ -78,7 +79,9 @@ class ClusterDescriptionModel(BaseClusterDescriptionModel): def __init__( self, - model: Union[str, "KnownModelName"] = "openai/gpt-4o-mini", + model: Union[ + str, "KnownModelName", "instructor.AsyncInstructor" + ] = "openai/gpt-4o-mini", max_concurrent_requests: int = 50, temperature: float = 0.2, checkpoint_filename: str = "clusters", @@ -88,20 +91,31 @@ def __init__( Initialize ClusterModel with core configuration. Args: - model: model identifier (e.g., "openai/gpt-4o-mini") + model: model identifier (e.g., "openai/gpt-4o-mini") or instructor client max_concurrent_requests: Maximum concurrent API requests temperature: LLM temperature for generation checkpoint_filename: Filename for checkpointing console: Rich console for progress tracking """ - self.model = model + import instructor + + # Handle both string model names and instructor client instances + if isinstance(model, str): + self.client = instructor.from_provider(model, async_client=True) + elif isinstance(model, instructor.AsyncInstructor): + self.client = model + else: + raise ValueError( + f"Invalid model type of type({type(model)}). Expected str or instructor.AsyncInstructor." + ) + self.max_concurrent_requests = max_concurrent_requests self.temperature = temperature self._checkpoint_filename = checkpoint_filename self.console = console logger.info( - f"Initialized ClusterModel with model={model}, max_concurrent_requests={max_concurrent_requests}, temperature={temperature}" + f"Initialized ClusterModel with client={self.client}, max_concurrent_requests={max_concurrent_requests}, temperature={temperature}" ) @property @@ -116,10 +130,8 @@ async def generate_clusters( max_contrastive_examples: int = 10, ) -> List[Cluster]: """Generate clusters from a mapping of cluster IDs to summaries.""" - import instructor self.sem = Semaphore(self.max_concurrent_requests) - self.client = instructor.from_provider(self.model, async_client=True) if not self.console: # Simple processing without rich display diff --git a/kura/meta_cluster.py b/kura/meta_cluster.py index 108ed37..d7a1ab3 100644 --- a/kura/meta_cluster.py +++ b/kura/meta_cluster.py @@ -16,11 +16,13 @@ from thefuzz import fuzz import asyncio import logging -from typing import Optional, Union +from typing import Optional, Union, TYPE_CHECKING -# Rich imports handled by Kura base class -from typing import TYPE_CHECKING +if TYPE_CHECKING: + from instructor.models import KnownModelName + import instructor +# Rich imports handled by Kura base class if TYPE_CHECKING: from rich.console import Console @@ -84,25 +86,36 @@ def checkpoint_filename(self) -> str: def __init__( self, + model: Union[ + str, "KnownModelName", "instructor.AsyncInstructor" + ] = "openai/gpt-4o-mini", max_concurrent_requests: int = 50, - model: str = "openai/gpt-4o-mini", embedding_model: Optional[BaseEmbeddingModel] = None, clustering_model: Union[BaseClusteringMethod, None] = None, max_clusters: int = 10, console: Optional["Console"] = None, **kwargs, # For future use ): + import instructor + if clustering_model is None: from kura.cluster import KmeansClusteringModel clustering_model = KmeansClusteringModel(12) + # Handle both string model names and instructor client instances + if isinstance(model, str): + self.client = instructor.from_provider(model, async_client=True) + elif isinstance(model, instructor.AsyncInstructor): + # Assume it's an instructor client + self.client = model + else: + raise ValueError( + f"Invalid model type of type({type(model)}). Expected str or instructor.AsyncInstructor." + ) + self.max_concurrent_requests = max_concurrent_requests self.sem = Semaphore(max_concurrent_requests) - - import instructor - - self.client = instructor.from_provider(model, async_client=True) self.console = console self.max_clusters = max_clusters @@ -111,11 +124,10 @@ def __init__( self.embedding_model = embedding_model self.clustering_model = clustering_model - self.model = model self.console = console logger.info( - f"Initialized MetaClusterModel with model={model}, max_concurrent_requests={max_concurrent_requests}, embedding_model={type(embedding_model).__name__}, clustering_model={type(clustering_model).__name__}, max_clusters={max_clusters}" + f"Initialized MetaClusterModel with client={self.client}, max_concurrent_requests={max_concurrent_requests}, embedding_model={type(embedding_model).__name__}, clustering_model={type(clustering_model).__name__}, max_clusters={max_clusters}" ) # Debug: Check if console is set diff --git a/kura/summarisation.py b/kura/summarisation.py index 7c99eba..d8571bb 100644 --- a/kura/summarisation.py +++ b/kura/summarisation.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from instructor.models import KnownModelName + import instructor from tqdm.asyncio import tqdm_asyncio from rich.console import Console @@ -98,7 +99,9 @@ class SummaryModel(BaseSummaryModel): def __init__( self, - model: Union[str, "KnownModelName"] = "openai/gpt-4o-mini", + model: Union[ + str, "KnownModelName", "instructor.AsyncInstructor" + ] = "openai/gpt-4o-mini", max_concurrent_requests: int = 50, checkpoint_filename: str = "summaries", console: Optional[Console] = None, @@ -110,11 +113,22 @@ def __init__( Per-use configuration (schemas, prompts, temperature) are method parameters. Args: - model: model identifier (e.g., "openai/gpt-4o-mini") + model: model identifier (e.g., "openai/gpt-4o-mini") or instructor client max_concurrent_requests: Maximum concurrent API requests cache: Caching strategy to use (optional) """ - self.model = model + import instructor + + # Handle both string model names and instructor client instances + if isinstance(model, str): + self.client = instructor.from_provider(model, async_client=True) + elif isinstance(model, instructor.AsyncInstructor): + self.client = model + else: + raise ValueError( + f"Invalid model type of type({type(model)}). Expected str or instructor.AsyncInstructor." + ) + self.max_concurrent_requests = max_concurrent_requests self._checkpoint_filename = checkpoint_filename self.console = console @@ -124,7 +138,7 @@ def __init__( cache_info = type(self.cache).__name__ if self.cache else "None" logger.info( - f"Initialized SummaryModel with model={model}, max_concurrent_requests={max_concurrent_requests}, cache={cache_info}" + f"Initialized SummaryModel with client={self.client}, max_concurrent_requests={max_concurrent_requests}, cache={cache_info}" ) @property @@ -201,20 +215,16 @@ async def summarise( self.semaphore = asyncio.Semaphore(self.max_concurrent_requests) logger.info( - f"Starting summarization of {len(conversations)} conversations using model {self.model}" + f"Starting summarization of {len(conversations)} conversations using client {self.client}" ) - import instructor - - client = instructor.from_provider(self.model, async_client=True) - if not self.console: # Simple progress tracking with tqdm summaries = await tqdm_asyncio.gather( *[ self._summarise_single_conversation( conversation, - client=client, + client=self.client, response_schema=response_schema, prompt=prompt, temperature=temperature, @@ -228,7 +238,7 @@ async def summarise( # Rich console progress tracking with live summary display summaries = await self._summarise_with_console( conversations, - client=client, + client=self.client, response_schema=response_schema, prompt=prompt, temperature=temperature, @@ -428,7 +438,7 @@ def update_preview_display(): for conversation in conversations: coro = self._summarise_single_conversation( conversation, - client=client, + client=self.client, response_schema=response_schema, prompt=prompt, temperature=temperature,