Model

Model wrappers for LLM and embedding APIs used across coshop.

Main exports:

• :class:LangChainModel — wraps a LangChain chat model with tool-calling support and a copy_for_prediction helper.
• :class:EmbeddingModelWrapper — wraps a SentenceTransformer or a remote embedding API with disk caching.
• Helper functions: :func:get_token_usage, :func:is_openai_model, :func:is_anthropic_model, :func:is_gemini_model.

EmbeddingModelWrapper(model_name, cache_dir=None, use_cache=True, max_cache_size_mb=None, max_cache_files=None, quantize=False, batch_size=128, cache_queries=False, device_map=None, embedding_api_url=None)

Wrapper for embedding models to handle different encoding methods. Supports standard SentenceTransformer models, Qwen3-Embedding, and EmbeddingGemma models.

Note: Qwen3-Embedding models require transformers>=4.51.0 and sentence-transformers>=2.7.0. EmbeddingGemma models work with standard sentence-transformers installation.

Includes disk caching for embeddings to avoid recomputing the same texts.

Initialize the embedding model wrapper.

Parameters:

Name	Type	Description	Default
model_name	str	Name of the model (e.g., "all-MiniLM-L6-v2", "Qwen/Qwen3-Embedding-0.6B", "google/embeddinggemma-300m")	required
cache_dir	Optional[str]	Optional directory for caching embeddings. If None, uses ".cache/embeddings/"	None
use_cache	bool	Whether to use disk caching for embeddings (default: True)	True
max_cache_size_mb	Optional[float]	Maximum cache size in MB. If None, defaults to 1024 MB (1 GB). Set to float('inf') for no limit.	None
max_cache_files	Optional[int]	Maximum number of cache files. If None, no file count limit (default: None)	None
quantize	Optional[bool]	Quantization level for the model. If None, no quantization is used.	False
batch_size	int	Batch size for encoding. Default is 128 (increased from SentenceTransformer's default of 32)	128
cache_queries	bool	Whether to cache query embeddings (default: False)	False
device_map	Optional[str]	Device map for multi-GPU (e.g. "auto"). Passed to SentenceTransformer via model_kwargs.	None
embedding_api_url	Optional[str]	If set, use OpenAI-compatible API (e.g. vLLM) instead of loading model locally.	None

encode(texts, show_progress_bar=False, **kwargs)

Encode texts into embeddings with caching. Handles different model types appropriately.

Parameters:

Name	Type	Description	Default
texts		List of texts or single text to encode	required
show_progress_bar	bool	Whether to show progress bar	False
**kwargs		Additional arguments passed to the encoding method	{}

Returns:

Type	Description
	numpy array of embeddings

encode_query(texts, show_progress_bar=False, **kwargs)

Encode queries into embeddings with caching. Uses model-specific query encoding when available.

Parameters:

Name	Type	Description	Default
texts		List of texts or single text to encode as queries	required
show_progress_bar	bool	Whether to show progress bar	False
**kwargs		Additional arguments passed to the encoding method	{}

Returns:

Type	Description
	numpy array of embeddings

LangChainModel(model_name, model_no_tools=None, summary_model_name=None, tools=None, verbosity=0, max_react_steps=25, min_react_steps=1, prompt_cache=True, multiturn_memory=False, summarize_state_after=None, out_of_steps_msg=None, list_tools_in_prompt=False, thinking_tokens=('<think>', '</think>'), add_thinking_tag=True, empty_message_filler=None, api_key=None, model_provider=None, summary_api_key=None, summary_model_provider=None, summary_vllm_api_url=None, **kwargs)

Bases: Model

raw_state property

Return the raw state of the chain (read-only) The raw-state only appends: no summarization, no deletion The exception is if multiturn_memory is False: then the raw_state will be [].

state property

Return the state of the chain

clear_state()

Clear the state of the chain

compress_state()

Compress the persisted chain by summarizing each tool result (see _compress_messages), then write the result back to the checkpointer. Final AI outputs are left untouched.

copy_for_prediction(*, tools, min_react_steps=None, max_react_steps=None)

Create a copy of this agent with different tools and/or react step limits. Preserves all other constructor arguments (model_provider, api_key, etc.). Note this new agent has an empty state. For prediction-time agents (e.g., agentic ranking or final recommendations), we force temperature=0.0 for more deterministic behavior, regardless of the interactive agent's temperature.

deduplicate_msgs_by_content(contents)

Delete all but the last occurrence of any message whose content is in contents.

get_state(slice=None)

Dump the state of the chain

insert_message(role, content, persist_state=True)

Insert a single message into the current conversation state.

Parameters:

Name	Type	Description	Default
role	str	One of "system", "user", "assistant", or "tool"	required
content	str	Message text content	required
persist_state	bool	If True, appends the message to the graph state; if False, no-op to state but returns the constructed message	True

Returns:

Type	Description
	The created LangChain BaseMessage

load_state(state)

Load the state of the chain

mark_history_cache_breakpoint()

Place an Anthropic prompt-cache breakpoint on the last message of the current state, so the whole conversation-history prefix is cached and re-read across subsequent generate() calls that share this history (e.g. prediction retries, the foregone-recall continuation, and the report/rank passes).

cache_control must live inside a content block, so a string-content message is converted to a single text block carrying the breakpoint, and a list-content message gets the breakpoint attached to its last text block. Returns False (no-op) for non-Anthropic models, when prompt caching is disabled, when the state is empty, or when no suitable text block is found. Stays within Anthropic's 4-breakpoint cap: it marks a single message and never accumulates across turns.

reset_token_tracking()

Zero the cumulative token counters (in place, to preserve any shared references held by prediction-time agent copies).

Model(name)

fmt_as_dialog(prompts=None, dialogs=None)

Parameters:

Name	Type	Description	Default
prompts	List[Union[str, Image]]	A list of prompts. The list dimension is over (B,) e.g. ["What is the capital of France?"] -> [[{"role": "user", "content": "What is the capital of France?"}]]	None
dialogs	List[List[Tuple[str, Union[str, Image]]]]	A list of [(role, content)] pairs. The list dimension is over (B, D) e.g. [[("user", "What is the capital of France?"), ("assistant", "Paris.")]] -> [[{"role": "user", "content": "What is the capital of France?"}, {"role": "assistant", "content": "Paris."}]]	None

Returns: Formatted dialogs: list of list of dictionaries (B, D) where each dictionary has keys "role" and "content"

generate(*, prompts=None, dialogs=None, **kwargs)

Parameters:

Name	Type	Description	Default
prompts	List[str]	A list of prompts. The list dimension is over (B,) e.g. ["What is the capital of France?"]	None
dialogs	List[List[Tuple[str, str]]]	A list of [(role, content)] pairs. The list dimension is over (B, D) e.g. [[("user", "What is the capital of France?"), ("assistant", "Paris.")]]	None

Returns: A list of completions. The list dimension is over (B,) e.g. ["Paris."]

encode_image_as_user_msg(image=None, image_path=None, extension='png', caption=None, model_name='gpt-5-nano')

Encode an image to a base64 string.

get_reasoning_effort_kwargs(model_name, reasoning_effort, max_tokens=64000)

Get the kwargs for the reasoning effort for a model.

get_token_breakdown(msg)

Extract per-message token usage broken down into four buckets.

Buckets

• input: total prompt tokens sent to the model (cached + uncached)
• input_cached: subset of input that was served from prompt cache
• output: total tokens generated by the model (includes reasoning for OpenAI)
• reasoning: subset of output tokens used for hidden reasoning (if any)

Prefers LangChain's normalized usage_metadata (uniform across OpenAI, Anthropic, Gemini). Falls back to provider-raw response_metadata when usage_metadata is unavailable. Returns zeros if neither is present (e.g. local HF/vLLM models that do not report usage).

get_token_usage(response_metadata, default_value=0, return_reasoning_tokens=True)

Get the token usage from the response metadata. If return_reasoning_tokens is True, return completion_tokens + reasoning_tokens. Otherwise, return just the completion tokens.

init_langchain_model(model_name, **kwargs)

Initialize a LangChain chat model. Args: model_name: The name of the model to initialize. **kwargs: Additional keyword arguments to pass to the model. Returns: A LangChain model.

is_anthropic_model(model_name, api_key=None) cached

Check if a model is an Anthropic model. Args: model_name: The name of the model to check. api_key: Optional API key for Anthropic client. Returns: True if the model is an Anthropic model, False otherwise.

is_openai_model(model_name, api_key=None) cached

Check if a model is an OpenAI model. Args: model_name: The name of the model to check. api_key: Optional API key for OpenAI client. Returns: True if the model is an OpenAI model, False otherwise.