471 lines
16 KiB
Python
471 lines
16 KiB
Python
|
"""Language Model (LLM) integration tool for OntoCast.
|
||
|
|
|
||
|
|
This module provides integration with various language models through LangChain,
|
||
|
|
supporting both OpenAI and Ollama providers. It enables text generation and
|
||
|
|
structured data extraction capabilities with optional caching support.
|
||
|
|
|
||
|
|
Cache Usage:
|
||
|
|
The LLM tool supports caching of responses to avoid redundant API calls.
|
||
|
|
Caching uses a shared Cacher instance that manages cache directories for all tools.
|
||
|
|
The cache directory is managed by the shared Cacher class and follows these rules:
|
||
|
|
|
||
|
|
```python
|
||
|
|
from ontocast.tool.llm import LLMTool
|
||
|
|
from ontocast.config import LLMConfig
|
||
|
|
from ontocast.tool.cache import Cacher
|
||
|
|
|
||
|
|
# Create shared cache instance
|
||
|
|
shared_cache = Cacher()
|
||
|
|
|
||
|
|
# Create LLM tool with shared cache
|
||
|
|
llm_tool = await LLMTool.acreate(
|
||
|
|
config=LLMConfig(...),
|
||
|
|
cache=shared_cache
|
||
|
|
)
|
||
|
|
```
|
||
|
|
|
||
|
|
Default cache locations:
|
||
|
|
- Tests: .test_cache/llm/ in the current working directory
|
||
|
|
- Windows: %USERPROFILE%\\AppData\\Local\\ontocast\\llm\
|
||
|
|
- Unix/Linux: ~/.cache/ontocast/llm/ (or $XDG_CACHE_HOME/ontocast/llm/)
|
||
|
|
|
||
|
|
Cache files are stored as JSON files with filenames based on SHA256 hashes
|
||
|
|
of the prompt and LLM configuration. This ensures that identical prompts
|
||
|
|
with the same configuration will return cached responses.
|
||
|
|
|
||
|
|
The shared Cacher automatically manages subdirectories for different tools,
|
||
|
|
ensuring organized cache storage while maintaining a single cache instance.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from functools import wraps
|
||
|
|
from typing import Any, Callable, Type, TypeVar
|
||
|
|
|
||
|
|
from langchain_core.language_models import BaseChatModel
|
||
|
|
from langchain_core.messages.ai import AIMessage
|
||
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
||
|
|
from langchain_ollama import ChatOllama
|
||
|
|
from langchain_openai import ChatOpenAI
|
||
|
|
from pydantic import BaseModel, Field, SecretStr
|
||
|
|
|
||
|
|
from ontocast.config import LLMConfig, LLMProvider
|
||
|
|
|
||
|
|
from .cache import Cacher, ToolCacher
|
||
|
|
from .onto import Tool
|
||
|
|
|
||
|
|
T = TypeVar("T", bound=BaseModel)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def track_llm_usage(func: Callable) -> Callable:
|
||
|
|
"""Decorator to track LLM usage automatically."""
|
||
|
|
|
||
|
|
@wraps(func)
|
||
|
|
def wrapper(self, *args, **kwargs):
|
||
|
|
# Get prompt for character counting
|
||
|
|
prompt = args[0] if args else ""
|
||
|
|
prompt_str = (
|
||
|
|
self._prompt_to_string(prompt)
|
||
|
|
if hasattr(self, "_prompt_to_string")
|
||
|
|
else str(prompt)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Call the original function
|
||
|
|
result = func(self, *args, **kwargs)
|
||
|
|
|
||
|
|
# Track usage if budget tracker is available in the tool
|
||
|
|
if hasattr(self, "budget_tracker") and self.budget_tracker is not None:
|
||
|
|
chars_sent = len(prompt_str)
|
||
|
|
chars_received = (
|
||
|
|
len(result.content)
|
||
|
|
if hasattr(result, "content") and result.content
|
||
|
|
else 0
|
||
|
|
)
|
||
|
|
self.budget_tracker.add_usage(chars_sent, chars_received)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
@wraps(func)
|
||
|
|
async def async_wrapper(self, *args, **kwargs):
|
||
|
|
# Get prompt for character counting
|
||
|
|
prompt = args[0] if args else ""
|
||
|
|
prompt_str = (
|
||
|
|
self._prompt_to_string(prompt)
|
||
|
|
if hasattr(self, "_prompt_to_string")
|
||
|
|
else str(prompt)
|
||
|
|
)
|
||
|
|
|
||
|
|
# Call the original function
|
||
|
|
result = await func(self, *args, **kwargs)
|
||
|
|
|
||
|
|
# Track usage if budget tracker is available in the tool
|
||
|
|
if hasattr(self, "budget_tracker") and self.budget_tracker is not None:
|
||
|
|
chars_sent = len(prompt_str)
|
||
|
|
chars_received = (
|
||
|
|
len(result.content)
|
||
|
|
if hasattr(result, "content") and result.content
|
||
|
|
else len(str(result))
|
||
|
|
)
|
||
|
|
self.budget_tracker.add_usage(chars_sent, chars_received)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
return async_wrapper if asyncio.iscoroutinefunction(func) else wrapper
|
||
|
|
|
||
|
|
|
||
|
|
class LLMTool(Tool):
|
||
|
|
"""Tool for interacting with language models.
|
||
|
|
|
||
|
|
This class provides a unified interface for working with different language model
|
||
|
|
providers (OpenAI, Ollama) through LangChain. It supports both synchronous and
|
||
|
|
asynchronous operations.
|
||
|
|
|
||
|
|
Attributes:
|
||
|
|
config: LLMConfig object containing all LLM settings.
|
||
|
|
cache: Cacher instance for caching LLM responses.
|
||
|
|
"""
|
||
|
|
|
||
|
|
config: LLMConfig = Field(default_factory=LLMConfig)
|
||
|
|
cache: Any = Field(default=None, exclude=True)
|
||
|
|
budget_tracker: Any = Field(default=None, exclude=True)
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
cache: Cacher | None = None,
|
||
|
|
budget_tracker: Any = None,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
"""Initialize the LLM tool.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
cache: Optional shared Cacher instance. If None, creates a new one.
|
||
|
|
budget_tracker: Optional budget tracker instance for usage statistics.
|
||
|
|
**kwargs: Additional keyword arguments passed to the parent class.
|
||
|
|
"""
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self._llm = None
|
||
|
|
self.budget_tracker = budget_tracker
|
||
|
|
|
||
|
|
# Initialize cache - use shared cacher or create new one
|
||
|
|
if cache is not None:
|
||
|
|
self.cache = ToolCacher(cache, "llm")
|
||
|
|
else:
|
||
|
|
# Fallback for backward compatibility
|
||
|
|
shared_cache = Cacher()
|
||
|
|
self.cache = ToolCacher(shared_cache, "llm")
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def create(
|
||
|
|
cls,
|
||
|
|
config: LLMConfig,
|
||
|
|
cache: Cacher | None = None,
|
||
|
|
budget_tracker: Any = None,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
"""Create a new LLM tool instance synchronously.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
config: LLMConfig object containing LLM settings.
|
||
|
|
cache: Optional shared Cacher instance.
|
||
|
|
budget_tracker: Optional budget tracker instance for usage statistics.
|
||
|
|
**kwargs: Additional keyword arguments for initialization.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
LLMTool: A new instance of the LLM tool.
|
||
|
|
"""
|
||
|
|
return asyncio.run(
|
||
|
|
cls.acreate(
|
||
|
|
config=config, cache=cache, budget_tracker=budget_tracker, **kwargs
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def acreate(
|
||
|
|
cls,
|
||
|
|
config: LLMConfig,
|
||
|
|
cache: Cacher | None = None,
|
||
|
|
budget_tracker: Any = None,
|
||
|
|
**kwargs,
|
||
|
|
):
|
||
|
|
"""Create a new LLM tool instance asynchronously.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
config: LLMConfig object containing LLM settings.
|
||
|
|
cache: Optional shared Cacher instance.
|
||
|
|
budget_tracker: Optional budget tracker instance for usage statistics.
|
||
|
|
**kwargs: Additional keyword arguments for initialization.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
LLMTool: A new instance of the LLM tool.
|
||
|
|
"""
|
||
|
|
# Create and initialize the instance with the config
|
||
|
|
self = cls(config=config, cache=cache, budget_tracker=budget_tracker, **kwargs)
|
||
|
|
await self.setup()
|
||
|
|
return self
|
||
|
|
|
||
|
|
async def setup(self):
|
||
|
|
"""Set up the language model based on the configured provider.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
ValueError: If the provider is not supported.
|
||
|
|
"""
|
||
|
|
if self.config.provider == LLMProvider.OPENAI:
|
||
|
|
if self.config.model_name.startswith("gpt-5"):
|
||
|
|
self.config.temperature = 1.0
|
||
|
|
logger.warning(
|
||
|
|
f"Setting temperature to {self.config.temperature} for gpt-5 class "
|
||
|
|
f"model {self.config.model_name}"
|
||
|
|
)
|
||
|
|
self._llm = ChatOpenAI(
|
||
|
|
model=self.config.model_name, # type: ignore
|
||
|
|
temperature=self.config.temperature,
|
||
|
|
base_url=self.config.base_url, # type: ignore
|
||
|
|
api_key=(
|
||
|
|
SecretStr(self.config.api_key) if self.config.api_key else None
|
||
|
|
), # type: ignore
|
||
|
|
)
|
||
|
|
elif self.config.provider == LLMProvider.OLLAMA:
|
||
|
|
self._llm = ChatOllama(
|
||
|
|
model=self.config.model_name,
|
||
|
|
base_url=self.config.base_url,
|
||
|
|
temperature=self.config.temperature,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise ValueError(f"Unsupported provider: {self.config.provider}")
|
||
|
|
|
||
|
|
@track_llm_usage
|
||
|
|
async def __call__(self, *args: Any, **kwds: Any) -> Any:
|
||
|
|
"""Call the language model directly (asynchronous).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
*args: Positional arguments passed to the LLM.
|
||
|
|
**kwds: Keyword arguments passed to the LLM.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Any: The LLM's response.
|
||
|
|
"""
|
||
|
|
# Extract prompt from args (first argument is typically the prompt)
|
||
|
|
prompt = args[0] if args else ""
|
||
|
|
|
||
|
|
# Prepare configuration for caching
|
||
|
|
config_dict = {
|
||
|
|
"provider": self.config.provider,
|
||
|
|
"model_name": self.config.model_name,
|
||
|
|
"temperature": self.config.temperature,
|
||
|
|
"base_url": self.config.base_url,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Check cache first
|
||
|
|
cached_response = self.cache.get(prompt, config=config_dict, **kwds)
|
||
|
|
|
||
|
|
if cached_response is not None:
|
||
|
|
prompt_str = self._prompt_to_string(prompt)
|
||
|
|
logger.debug(f"Cache hit for __call__: {prompt_str[:50]}...")
|
||
|
|
# Return a mock BaseMessage object with the cached content
|
||
|
|
content = cached_response["content"]
|
||
|
|
content_str = content if isinstance(content, str) else str(content)
|
||
|
|
return AIMessage(content=content_str)
|
||
|
|
|
||
|
|
# Generate new response
|
||
|
|
prompt_str = self._prompt_to_string(prompt)
|
||
|
|
logger.debug(
|
||
|
|
f"Cache miss, calling LLM for __call__, prompt size {len(prompt_str[:50])}..."
|
||
|
|
)
|
||
|
|
|
||
|
|
response = await self.llm.ainvoke(*args, **kwds)
|
||
|
|
|
||
|
|
# Cache the response
|
||
|
|
response_data = {
|
||
|
|
"content": response.content,
|
||
|
|
"prompt": self._prompt_to_string(prompt),
|
||
|
|
"kwargs": kwds,
|
||
|
|
}
|
||
|
|
self.cache.set(prompt, response_data, config=config_dict, **kwds)
|
||
|
|
|
||
|
|
return response
|
||
|
|
|
||
|
|
@track_llm_usage
|
||
|
|
async def acall(self, *args: Any, **kwds: Any) -> Any:
|
||
|
|
"""Call the language model directly (asynchronous).
|
||
|
|
|
||
|
|
Args:
|
||
|
|
*args: Positional arguments passed to the LLM.
|
||
|
|
**kwds: Keyword arguments passed to the LLM.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Any: The LLM's response.
|
||
|
|
"""
|
||
|
|
# Extract prompt from args (first argument is typically the prompt)
|
||
|
|
prompt = args[0] if args else ""
|
||
|
|
|
||
|
|
# Prepare configuration for caching
|
||
|
|
config_dict = {
|
||
|
|
"provider": self.config.provider,
|
||
|
|
"model_name": self.config.model_name,
|
||
|
|
"temperature": self.config.temperature,
|
||
|
|
"base_url": self.config.base_url,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Check cache first
|
||
|
|
cached_response = self.cache.get(prompt, config=config_dict, **kwds)
|
||
|
|
|
||
|
|
if cached_response is not None:
|
||
|
|
prompt_str = self._prompt_to_string(prompt)
|
||
|
|
logger.debug(f"Cache hit for acall: {prompt_str[:50]}...")
|
||
|
|
# Return a mock BaseMessage object with the cached content
|
||
|
|
content = cached_response["content"]
|
||
|
|
content_str = content if isinstance(content, str) else str(content)
|
||
|
|
return AIMessage(content=content_str)
|
||
|
|
|
||
|
|
# Generate new response
|
||
|
|
prompt_str = self._prompt_to_string(prompt)
|
||
|
|
logger.debug(f"Cache miss, calling LLM for acall: {prompt_str[:50]}...")
|
||
|
|
|
||
|
|
response = await self.llm.ainvoke(*args, **kwds)
|
||
|
|
|
||
|
|
# Cache the response
|
||
|
|
response_data = {
|
||
|
|
"content": response.content,
|
||
|
|
"prompt": self._prompt_to_string(prompt),
|
||
|
|
"kwargs": kwds,
|
||
|
|
}
|
||
|
|
self.cache.set(prompt, response_data, config=config_dict, **kwds)
|
||
|
|
|
||
|
|
return response
|
||
|
|
|
||
|
|
@property
|
||
|
|
def llm(self) -> BaseChatModel:
|
||
|
|
"""Get the underlying language model instance.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
BaseChatModel: The configured language model.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
RuntimeError: If the LLM has not been properly initialized.
|
||
|
|
"""
|
||
|
|
if self._llm is None:
|
||
|
|
raise RuntimeError(
|
||
|
|
"LLM resource not properly initialized. Call setup() first."
|
||
|
|
)
|
||
|
|
return self._llm
|
||
|
|
|
||
|
|
def _prompt_to_string(self, prompt) -> str:
|
||
|
|
"""Convert various prompt types to string for caching.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt: The prompt object (string, StringPromptValue, etc.)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
str: String representation of the prompt.
|
||
|
|
"""
|
||
|
|
if isinstance(prompt, str):
|
||
|
|
return prompt
|
||
|
|
elif hasattr(prompt, "to_string"):
|
||
|
|
return prompt.to_string()
|
||
|
|
elif hasattr(prompt, "text"):
|
||
|
|
return prompt.text
|
||
|
|
elif hasattr(prompt, "content"):
|
||
|
|
return prompt.content
|
||
|
|
else:
|
||
|
|
return str(prompt)
|
||
|
|
|
||
|
|
@track_llm_usage
|
||
|
|
async def complete(self, prompt: str, **kwargs) -> Any:
|
||
|
|
"""Generate a completion for the given prompt.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt: The input prompt for generation.
|
||
|
|
**kwargs: Additional keyword arguments for generation.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Any: The generated completion.
|
||
|
|
"""
|
||
|
|
# Prepare configuration for caching
|
||
|
|
config_dict = {
|
||
|
|
"provider": self.config.provider,
|
||
|
|
"model_name": self.config.model_name,
|
||
|
|
"temperature": self.config.temperature,
|
||
|
|
"base_url": self.config.base_url,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Check cache first
|
||
|
|
cached_response = self.cache.get(prompt, config=config_dict, **kwargs)
|
||
|
|
|
||
|
|
if cached_response is not None:
|
||
|
|
logger.debug(f"Cache hit for prompt: {prompt[:50]}...")
|
||
|
|
content = cached_response["content"]
|
||
|
|
return content if isinstance(content, str) else str(content)
|
||
|
|
|
||
|
|
# Generate new response
|
||
|
|
logger.debug(f"Cache miss, calling LLM for prompt: {prompt[:50]}...")
|
||
|
|
|
||
|
|
response = await self.llm.ainvoke(prompt, **kwargs)
|
||
|
|
|
||
|
|
# Cache the response
|
||
|
|
response_data = {
|
||
|
|
"content": response.content,
|
||
|
|
"prompt": self._prompt_to_string(prompt),
|
||
|
|
"kwargs": kwargs,
|
||
|
|
}
|
||
|
|
self.cache.set(prompt, response_data, config=config_dict, **kwargs)
|
||
|
|
|
||
|
|
return response.content
|
||
|
|
|
||
|
|
@track_llm_usage
|
||
|
|
async def extract(self, prompt: str, output_schema: Type[T], **kwargs) -> T:
|
||
|
|
"""Extract structured data from the prompt according to a schema.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
prompt: The input prompt for extraction.
|
||
|
|
output_schema: The Pydantic model class defining the output structure.
|
||
|
|
**kwargs: Additional keyword arguments for extraction.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
T: The extracted data conforming to the output schema.
|
||
|
|
"""
|
||
|
|
parser = PydanticOutputParser(pydantic_object=output_schema)
|
||
|
|
format_instructions = parser.get_format_instructions()
|
||
|
|
|
||
|
|
full_prompt = f"{prompt}\n\n{format_instructions}"
|
||
|
|
|
||
|
|
# Prepare configuration for caching
|
||
|
|
config_dict = {
|
||
|
|
"provider": self.config.provider,
|
||
|
|
"model_name": self.config.model_name,
|
||
|
|
"temperature": self.config.temperature,
|
||
|
|
"base_url": self.config.base_url,
|
||
|
|
"output_schema": output_schema.__name__,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Check cache first
|
||
|
|
cached_response = self.cache.get(full_prompt, config=config_dict, **kwargs)
|
||
|
|
|
||
|
|
if cached_response is not None:
|
||
|
|
logger.debug(f"Cache hit for extraction: {prompt[:50]}...")
|
||
|
|
# Parse the cached content
|
||
|
|
content = cached_response["content"]
|
||
|
|
if isinstance(content, str):
|
||
|
|
return parser.parse(content)
|
||
|
|
else:
|
||
|
|
# Fallback: convert to string if it's not already
|
||
|
|
return parser.parse(str(content))
|
||
|
|
|
||
|
|
# Generate new response
|
||
|
|
logger.debug(f"Cache miss, calling LLM for extraction: {prompt[:50]}...")
|
||
|
|
|
||
|
|
response = await self.llm.ainvoke(full_prompt, **kwargs)
|
||
|
|
|
||
|
|
# Cache the response
|
||
|
|
response_data = {
|
||
|
|
"content": response.content,
|
||
|
|
"prompt": self._prompt_to_string(full_prompt),
|
||
|
|
"output_schema": output_schema.__name__,
|
||
|
|
"kwargs": kwargs,
|
||
|
|
}
|
||
|
|
self.cache.set(full_prompt, response_data, config=config_dict, **kwargs)
|
||
|
|
|
||
|
|
content = response.content
|
||
|
|
return parser.parse(content if isinstance(content, str) else str(content))
|