portainer-core/services/core-ai/src/agent.py

"""
Core AI Agent - Direct LiteLLM Chat Completion
This is a diagnostic file to test direct text generation via LiteLLM, bypassing Google ADK.
"""
import os
import logging
from typing import AsyncIterator, Dict, Any, List, Optional
from functools import lru_cache

# We will directly use litellm here
import litellm

# Adjusted import paths for the new core-ai service structure
from src.config import get_settings
from src.prompts import get_prompt

logger = logging.getLogger(__name__)


# Simplified Agent for direct LiteLLM interaction
class SimpleLiteLLMAgent:
    def __init__(self):
        # Enable verbose logging for LiteLLM
        litellm.set_verbose = True
        logger.info("LiteLLM verbose logging enabled.")

        self.settings = get_settings()

        # Load system prompt
        self.system_prompt = get_prompt(self.settings.system_prompt_variant)
        logger.info(f"System prompt variant: {self.settings.system_prompt_variant}")
        logger.info(f"System prompt: {self.system_prompt[:100]}...")

        # Initialize LiteLLM for Ollama (format: "ollama/model_name")
        model_name = self.settings.agent_model
        litellm_model = f"ollama/{model_name}"

        logger.info(f"Initializing LiteLLM direct model: {litellm_model}")
        logger.info(f"Ollama base URL from settings: {self.settings.ollama_base_url}")

        self.model_params = {
            "model": litellm_model,
            "api_base": self.settings.ollama_base_url,
            "temperature": 0.1,
            # No tool definitions passed here to force text generation
        }

    async def chat(
        self,
        messages: List[Dict[str, str]],
        conversation_id: str = None, # Not used in this simple mode
        stream: bool = True,
        prompt_variant: Optional[str] = None # Not used in this simple mode
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        Processes a chat message using direct LiteLLM completion.
        """
        logger.info(f"🚀 Starting direct LiteLLM completion for message: {messages[-1]['content'][:50]}...")
        try:
            # Prepare messages in LiteLLM format
            litellm_messages = [{"role": m["role"], "content": m["content"]} for m in messages]

            # Inject system prompt if not already present
            if not litellm_messages or litellm_messages[0]["role"] != "system":
                litellm_messages.insert(0, {"role": "system", "content": self.system_prompt})
                logger.info("✓ System prompt injected")

            # Log full message payload for debugging
            logger.info(f"📤 Sending {len(litellm_messages)} messages to LiteLLM:")
            for i, msg in enumerate(litellm_messages):
                content_preview = msg['content'][:100] + "..." if len(msg['content']) > 100 else msg['content']
                logger.info(f"  [{i}] {msg['role']}: {content_preview}")

            # Use acompletion for async environments
            response = await litellm.acompletion(
                messages=litellm_messages,
                stream=stream,
                **self.model_params
            )

            if stream:
                chunk_count = 0
                async for chunk in response:
                    chunk_count += 1
                    content_delta = chunk.choices[0].delta.content if chunk.choices[0].delta.content else ""
                    finish_reason = chunk.choices[0].finish_reason
                    if content_delta:
                        yield {"type": "content", "content": content_delta}
                    if finish_reason:
                        logger.info(f"📥 Stream completed after {chunk_count} chunks. Finish reason: {finish_reason}")
                        yield {"type": "content", "content": "", "finish_reason": finish_reason}
            else:
                content = response.choices[0].message.content
                logger.info(f"📥 Response received: {content[:200]}..." if len(content) > 200 else f"📥 Response received: {content}")
                yield {"type": "content", "content": content, "finish_reason": "stop"}

        except Exception as e:
            logger.error(f"Error in direct LiteLLM chat: {e}", exc_info=True)
            yield {
                "type": "error",
                "content": f"Sorry, an error occurred during text generation: {str(e)}",
                "finish_reason": "stop"
            }

    async def chat_completion(
        self,
        messages: List[Dict[str, str]],
        conversation_id: str = None,
        prompt_variant: Optional[str] = None
    ) -> str:
        """
        Get a non-streaming response from the direct LiteLLM chat.
        """
        final_content = ""
        async for chunk in self.chat(messages=messages, conversation_id=conversation_id, stream=False, prompt_variant=prompt_variant):
            if chunk["type"] == "content":
                final_content += chunk["content"]
                if chunk.get("finish_reason") == "stop":
                    break
        return final_content if final_content else "I couldn't generate a response."


@lru_cache()
def get_simple_litellm_agent() -> SimpleLiteLLMAgent:
    """Get cached simple LiteLLM agent instance"""
    return SimpleLiteLLMAgent()