93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
|
"""LLM Adapter
|
||
|
|
|
||
|
|
Thin abstraction over a concrete LLM backend (local or remote).
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from typing import Any, Callable, Dict, Optional
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class LLMResponse:
|
||
|
|
"""Container for a single LLM interaction."""
|
||
|
|
|
||
|
|
prompt: str
|
||
|
|
response: str
|
||
|
|
tokens_prompt: Optional[int] = None
|
||
|
|
tokens_response: Optional[int] = None
|
||
|
|
model_name: Optional[str] = None
|
||
|
|
metadata: Dict[str, Any] = None
|
||
|
|
|
||
|
|
def to_dict(self) -> Dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"prompt": self.prompt,
|
||
|
|
"response": self.response,
|
||
|
|
"tokens_prompt": self.tokens_prompt,
|
||
|
|
"tokens_response": self.tokens_response,
|
||
|
|
"model_name": self.model_name,
|
||
|
|
"metadata": self.metadata or {},
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
class LLMAdapter:
|
||
|
|
"""Adapter around a concrete LLM backend.
|
||
|
|
|
||
|
|
backend: Callable that takes (prompt, **kwargs) and returns:
|
||
|
|
- str
|
||
|
|
- or dict with keys like:
|
||
|
|
- response
|
||
|
|
- tokens_prompt
|
||
|
|
- tokens_response
|
||
|
|
- model_name
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
backend: Callable[..., Any],
|
||
|
|
model_name: Optional[str] = None,
|
||
|
|
) -> None:
|
||
|
|
self._backend = backend
|
||
|
|
self._model_name = model_name or "unknown"
|
||
|
|
|
||
|
|
def run(self, prompt: str, **kwargs: Any) -> LLMResponse:
|
||
|
|
"""Run the underlying LLM on a prompt and normalize the result."""
|
||
|
|
raw = self._backend(prompt, **kwargs)
|
||
|
|
|
||
|
|
if isinstance(raw, str):
|
||
|
|
return LLMResponse(
|
||
|
|
prompt=prompt,
|
||
|
|
response=raw,
|
||
|
|
model_name=self._model_name,
|
||
|
|
metadata={},
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(raw, dict):
|
||
|
|
return LLMResponse(
|
||
|
|
prompt=prompt,
|
||
|
|
response=str(raw.get("response", "")),
|
||
|
|
tokens_prompt=raw.get("tokens_prompt"),
|
||
|
|
tokens_response=raw.get("tokens_response"),
|
||
|
|
model_name=raw.get("model_name", self._model_name),
|
||
|
|
metadata={
|
||
|
|
k: v
|
||
|
|
for k, v in raw.items()
|
||
|
|
if k
|
||
|
|
not in {
|
||
|
|
"response",
|
||
|
|
"tokens_prompt",
|
||
|
|
"tokens_response",
|
||
|
|
"model_name",
|
||
|
|
}
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
# Fallback: best-effort stringification
|
||
|
|
return LLMResponse(
|
||
|
|
prompt=prompt,
|
||
|
|
response=str(raw),
|
||
|
|
model_name=self._model_name,
|
||
|
|
metadata={"raw_type": type(raw).__name__},
|
||
|
|
)
|