File size: 2,648 Bytes
3bc3e37
68ee171
 
be70bd4
3bc3e37
68ee171
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3bc3e37
00b17e7
68ee171
3bc3e37
484a699
00b17e7
 
68ee171
00b17e7
 
 
484a699
 
 
68ee171
00b17e7
 
 
be70bd4
 
68ee171
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
be70bd4
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional
import re


@dataclass
class ToolCallInfo:
    """Record of a single tool invocation."""

    name: str
    args: dict = field(default_factory=dict)


@dataclass
class ToolCallResult:
    """Structured return value from ``generate_with_tools``."""

    text: str
    used_tool: bool
    tool_name: Optional[str] = None
    tool_args: dict = field(default_factory=dict)
    tool_calls_made: List["ToolCallInfo"] = field(default_factory=list)


class LLMClient(ABC):
    """Abstract base class for all LLM clients"""

    @abstractmethod
    async def generate(self, system_prompt: str, context: List[dict], temperature: float, max_tokens: int, response_mime_type: str = None) -> str:
        """
        Generate a response using the LLM.

        Args:
            system_prompt (str): The system prompt defining the persona/role
            context (List[dict]): List of conversation messages with 'role' and 'content' keys
            temperature (float): Sampling temperature for generation
            max_tokens (int): Maximum number of tokens to generate
            response_mime_type (str, optional): MIME type for the response format. Defaults to None.

        Returns:
            str: The generated response text
        """
        pass

    async def generate_with_tools(
        self,
        system_prompt: str,
        user_message: str,
        tool_definitions: Optional[List[Dict[str, Any]]] = None,
        tool_executor: Optional[Callable] = None,
        temperature: float = 0.7,
        max_tokens: int = 2048,
    ) -> ToolCallResult:
        """Generate a response, optionally invoking tools.

        Subclasses that support native tool calling should override this
        method.  The default implementation ignores tools and falls back
        to a plain ``generate()`` call so that providers without tool
        support degrade gracefully.
        """
        text = await self.generate(
            system_prompt=system_prompt,
            context=[{"role": "user", "content": user_message}],
            temperature=temperature,
            max_tokens=max_tokens,
        )
        return ToolCallResult(text=text, used_tool=False)

    def _clean_response(self, response: str) -> str:
        """Clean up response text, preserving Markdown formatting."""
        response = response.replace("\r\n", "\n").replace("\r", "\n")
        lines = [ln.rstrip() for ln in response.split("\n")]
        response = re.sub(r"\n{3,}", "\n\n", "\n".join(lines)).strip()
        return response