From 1d57d8bc8a0abf4772723a09eaf4b79016ea9fc4 Mon Sep 17 00:00:00 2001 From: Philipp Temminghoff Date: Tue, 19 Nov 2024 18:20:55 +0100 Subject: [PATCH] chore: cleanup --- src/llmling/client.py | 6 ++++-- src/llmling/llm/providers/litellm.py | 5 ++--- .../processors/implementations/template.py | 2 +- src/llmling/resources/web_research.yml | 2 +- src/llmling/task/executor.py | 16 ++++++++++------ src/llmling/task/manager.py | 16 +++++++++++++--- 6 files changed, 31 insertions(+), 16 deletions(-) diff --git a/src/llmling/client.py b/src/llmling/client.py index 685ee80..7fa6d54 100644 --- a/src/llmling/client.py +++ b/src/llmling/client.py @@ -146,7 +146,7 @@ async def startup(self) -> None: # Start processor registry await self._processor_registry.startup() - # Create executor and manager + # Create executor with empty tool registry self._executor = TaskExecutor( context_registry=context_registry, processor_registry=self._processor_registry, @@ -154,6 +154,8 @@ async def startup(self) -> None: tool_registry=self.tool_registry, config_manager=self.config_manager, ) + + # Create manager (will handle tool registration) self._manager = TaskManager(self.config_manager.config, self._executor) self._initialized = True @@ -161,7 +163,7 @@ async def startup(self) -> None: except Exception as exc: logger.exception("Client initialization failed") - await self.shutdown() # Ensure cleanup on failure + await self.shutdown() msg = f"Failed to initialize client: {exc}" raise exceptions.LLMLingError(msg) from exc diff --git a/src/llmling/llm/providers/litellm.py b/src/llmling/llm/providers/litellm.py index b690de7..1d3b1d0 100644 --- a/src/llmling/llm/providers/litellm.py +++ b/src/llmling/llm/providers/litellm.py @@ -141,9 +141,8 @@ async def _complete_impl( ) except Exception as exc: - logger.exception("LiteLLM completion failed with error:") - error_msg = f"LiteLLM completion failed: {exc}" - raise exceptions.LLMError(error_msg) from exc + error_msg = f"Failed to process LiteLLM response: {exc}" + raise ValueError(error_msg) from exc def _is_local_provider(self) -> bool: """Check if the current model is a local provider (like Ollama).""" diff --git a/src/llmling/processors/implementations/template.py b/src/llmling/processors/implementations/template.py index c3fb9c4..cf8c38d 100644 --- a/src/llmling/processors/implementations/template.py +++ b/src/llmling/processors/implementations/template.py @@ -62,7 +62,7 @@ async def _process_impl(self, context: ProcessingContext) -> ProcessorResult: original_content=context.original_content, metadata={ "template_vars": list(render_context.keys()), - "template": self.config.template, # Add this for debugging + "template": self.config.template, }, ) except Exception as exc: diff --git a/src/llmling/resources/web_research.yml b/src/llmling/resources/web_research.yml index 89aa1ea..fad93b1 100644 --- a/src/llmling/resources/web_research.yml +++ b/src/llmling/resources/web_research.yml @@ -6,7 +6,7 @@ tools: llm_providers: llama: name: "Local Llama" - model: "ollama/llama3.2:3b" + model: "openai/gpt-3.5-turbo" contexts: research_prompt: diff --git a/src/llmling/task/executor.py b/src/llmling/task/executor.py index c40ed31..5e39006 100644 --- a/src/llmling/task/executor.py +++ b/src/llmling/task/executor.py @@ -34,7 +34,7 @@ def __init__( processor_registry: ProcessorRegistry, provider_registry: ProviderRegistry, tool_registry: ToolRegistry | None = None, - config_manager: ConfigManager | None = None, # Add this + config_manager: ConfigManager | None = None, *, default_timeout: int = 30, default_max_retries: int = 3, @@ -97,9 +97,15 @@ def _prepare_tool_config( # Get schemas for all available tools tool_schemas: list[dict[str, Any]] = [] for tool_name in available_tools: + # Verify tool exists before getting schema + if not self.tool_registry.has_tool(tool_name): + logger.warning("Tool not found in registry: %s", tool_name) + continue + schema = self.tool_registry.get_schema(tool_name) logger.debug("Tool schema for %s: %s", tool_name, schema) tool_schemas.append(schema.function) + return { "tools": tool_schemas, "tool_choice": ( @@ -125,6 +131,7 @@ async def execute( try: # Add tool configuration if available if tool_config := self._prepare_tool_config(task_context, task_provider): + logger.debug("Tool configuration prepared: %s", tool_config) kwargs.update(tool_config) # Load and process context context_result = await self._load_context(task_context) @@ -150,11 +157,8 @@ async def execute( if completion.tool_calls: tool_results = [] for tool_call in completion.tool_calls: - logger.debug( - "Executing tool call: %s with params: %s", - tool_call.name, - tool_call.parameters, - ) # Add this + msg = "Executing tool call: %s with params: %s" + logger.debug(msg, tool_call.name, tool_call.parameters) result = await self.tool_registry.execute( tool_call.name, **tool_call.parameters, diff --git a/src/llmling/task/manager.py b/src/llmling/task/manager.py index f96969a..996cb06 100644 --- a/src/llmling/task/manager.py +++ b/src/llmling/task/manager.py @@ -30,8 +30,6 @@ def __init__( """Initialize task manager.""" self.config = config self.executor = executor - - # Register tools only if they exist in config if self.config.tools: self._register_tools() @@ -86,16 +84,28 @@ async def execute_template( **kwargs: Any, ) -> TaskResult: """Execute a task template.""" + # First do all validation/setup without retry try: task_context, task_provider = self._prepare_task(template_name, system_prompt) + except exceptions.TaskError: + # Configuration errors should fail immediately + logger.exception("Task preparation failed") + raise + + # Then execute with retry only for LLM-related errors + try: return await self.executor.execute( task_context, task_provider, system_prompt=system_prompt, **kwargs, ) - except Exception as exc: + except exceptions.LLMError: + # LLM errors can be retried by the provider logger.exception("Task execution failed") + raise + except Exception as exc: + # Other errors should not be retried msg = f"Task execution failed for template {template_name}: {exc}" raise exceptions.TaskError(msg) from exc