Skip to content

Flexible model client #163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ivanleomk wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Flexible model client #163

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
211431e
fix: migrate cluster to use instructor client or the model naem
ivanleomk Jun 30, 2025
230c7b1
fix: adding supp for summary and meta cluster
ivanleomk Jun 30, 2025
6b8584f
fix:
ivanleomk Jun 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions kura/cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
if TYPE_CHECKING:
from instructor.models import KnownModelName
from instructor import AsyncInstructor
import instructor

import numpy as np
import asyncio
Expand Down Expand Up @@ -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",
Expand All @@ -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)

@ellipsis-dev ellipsis-dev Bot Jun 30, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Runtime error: 'instructor' is referenced in init but only imported inside a TYPE_CHECKING block. Import it at runtime (e.g. add a local import) to avoid NameError.

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
Expand All @@ -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
Expand Down
32 changes: 22 additions & 10 deletions kura/meta_cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down
34 changes: 22 additions & 12 deletions kura/summarisation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down