db70fdbe5e
* feat: add LiteLLM LLM provider for Bedrock and 100+ providers Add a new `litellm` LLM provider that uses the LiteLLM SDK for chat completions and tool calling, enabling AWS Bedrock and 100+ other providers for Hindsight's core engine (retain, recall, reflect). - New LiteLLMLLM provider in engine/providers/litellm_llm.py - Registered in factory, valid providers list, and no-api-key set - Refactored API key validation to use requires_api_key() helper - Added boto3 dependency for Bedrock auth - Updated docs: configuration, models, monitoring, providers grid * feat: add bedrock as first-class LLM provider alias Add `bedrock` as a dedicated provider name that auto-prepends the `bedrock/` prefix to model names and delegates to LiteLLMLLM under the hood. This makes Bedrock support more discoverable — users set `HINDSIGHT_API_LLM_PROVIDER=bedrock` with plain Bedrock model IDs. * test: add Bedrock to CI provider tests - Add bedrock/us.amazon.nova-lite-v1:0 to MODEL_MATRIX in test_llm_provider.py - Add AWS credential check in should_skip_provider() - Pass AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION_NAME secrets to test-api job - Update default bedrock model to amazon.nova-2-lite-v1:0 * fix: regenerate docs skill files and bump memory test timeout - Regenerate skills/hindsight-docs references after docs changes - Bump test_llm_provider_memory_operations timeout to 600s for slower providers like Bedrock via LiteLLM * test: skip bedrock lite models in memory operations test Nova Lite has a 10K output token limit which is too low for fact extraction (requires 64K). The api_methods test (completion, tools, structured output) already validates the provider works correctly. * test: use Nova Pro for bedrock CI tests to cover full memory pipeline Nova Lite only supports 10K output tokens, too low for fact extraction. Switch to Nova Pro which supports the full 64K output needed for retain/reflect operations. This ensures bedrock is tested on all Hindsight functionalities, not just basic API methods. * test: switch bedrock CI to Nova 2 Lite (supports 64K output tokens) Nova v1 models (Pro, Lite) have a 10K output token limit which is too low for fact extraction. Nova 2 Lite supports 64K+ output tokens, enabling full memory pipeline testing (retain + reflect).
380 lines
14 KiB
Python
380 lines
14 KiB
Python
"""
|
|
LiteLLM LLM provider for universal model support.
|
|
|
|
This provider enables using 100+ LLM providers via the LiteLLM SDK, including:
|
|
- AWS Bedrock (bedrock/anthropic.claude-3-5-sonnet-...)
|
|
- Azure OpenAI (azure/gpt-4o)
|
|
- Together AI (together_ai/meta-llama/...)
|
|
- Any other LiteLLM-supported provider
|
|
|
|
Uses litellm.acompletion() for async chat completions.
|
|
Authentication for cloud providers (e.g., AWS Bedrock via boto3 credential chain)
|
|
is handled automatically by LiteLLM.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
|
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
|
from hindsight_api.metrics import get_metrics_collector
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LiteLLMLLM(LLMInterface):
|
|
"""
|
|
LLM provider using the LiteLLM SDK for universal model support.
|
|
|
|
Supports any model accessible via litellm.acompletion(), including AWS Bedrock,
|
|
Azure OpenAI, Together AI, Fireworks AI, and more.
|
|
|
|
Model names follow LiteLLM conventions with provider prefixes:
|
|
- bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
|
|
- azure/gpt-4o
|
|
- together_ai/meta-llama/Llama-3-70b-chat-hf
|
|
- fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
provider: str,
|
|
api_key: str,
|
|
base_url: str,
|
|
model: str,
|
|
reasoning_effort: str = "low",
|
|
timeout: float = 300.0,
|
|
**kwargs: Any,
|
|
):
|
|
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
|
self.timeout = timeout
|
|
self._litellm: Any = None
|
|
|
|
try:
|
|
import litellm
|
|
|
|
self._litellm = litellm
|
|
# Suppress LiteLLM's verbose logging
|
|
litellm.suppress_debug_info = True # type: ignore[assignment]
|
|
# Drop unsupported params instead of raising errors (e.g. tool_choice on some Bedrock models)
|
|
litellm.drop_params = True # type: ignore[assignment]
|
|
logging.getLogger("LiteLLM").setLevel(logging.WARNING)
|
|
logger.info(f"LiteLLM SDK initialized for model: {self.model}")
|
|
except ImportError as e:
|
|
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
|
|
|
|
async def verify_connection(self) -> None:
|
|
try:
|
|
test_messages = [{"role": "user", "content": "test"}]
|
|
await self.call(
|
|
messages=test_messages,
|
|
max_completion_tokens=50,
|
|
temperature=0.0,
|
|
scope="verification",
|
|
max_retries=0,
|
|
)
|
|
logger.info("LiteLLM connection verified successfully")
|
|
except OutputTooLongError:
|
|
# Truncation is fine for verification — it means the connection works
|
|
logger.info("LiteLLM connection verified successfully (response truncated)")
|
|
except Exception as e:
|
|
logger.error(f"LiteLLM connection verification failed: {e}")
|
|
raise RuntimeError(f"Failed to verify LiteLLM connection: {e}") from e
|
|
|
|
def _build_common_kwargs(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
max_completion_tokens: int | None = None,
|
|
temperature: float | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Build common kwargs for litellm calls."""
|
|
kwargs: dict[str, Any] = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"timeout": self.timeout,
|
|
}
|
|
|
|
if self.api_key:
|
|
kwargs["api_key"] = self.api_key
|
|
if self.base_url:
|
|
kwargs["api_base"] = self.base_url
|
|
if max_completion_tokens is not None:
|
|
kwargs["max_completion_tokens"] = max_completion_tokens
|
|
if temperature is not None:
|
|
kwargs["temperature"] = temperature
|
|
|
|
return kwargs
|
|
|
|
async def call(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
response_format: Any | None = None,
|
|
max_completion_tokens: int | None = None,
|
|
temperature: float | None = None,
|
|
scope: str = "memory",
|
|
max_retries: int = 10,
|
|
initial_backoff: float = 1.0,
|
|
max_backoff: float = 60.0,
|
|
skip_validation: bool = False,
|
|
strict_schema: bool = False,
|
|
return_usage: bool = False,
|
|
) -> Any:
|
|
start_time = time.time()
|
|
|
|
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
|
|
|
# Add JSON schema response format if provided
|
|
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
|
schema = response_format.model_json_schema()
|
|
call_kwargs["response_format"] = {
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": response_format.__name__ if hasattr(response_format, "__name__") else "response",
|
|
"schema": schema,
|
|
"strict": strict_schema,
|
|
},
|
|
}
|
|
|
|
last_exception = None
|
|
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
response = await self._litellm.acompletion(**call_kwargs)
|
|
|
|
content = response.choices[0].message.content or ""
|
|
finish_reason = response.choices[0].finish_reason
|
|
|
|
# Check for length-limited output
|
|
if finish_reason == "length":
|
|
raise OutputTooLongError("LiteLLM response was truncated due to token limit")
|
|
|
|
if response_format is not None:
|
|
# Strip markdown code fences if present
|
|
clean_content = content
|
|
if "```json" in content:
|
|
clean_content = content.split("```json")[1].split("```")[0].strip()
|
|
elif "```" in content:
|
|
clean_content = content.split("```")[1].split("```")[0].strip()
|
|
|
|
try:
|
|
json_data = json.loads(clean_content)
|
|
except json.JSONDecodeError:
|
|
json_data = json.loads(content)
|
|
|
|
if skip_validation:
|
|
result = json_data
|
|
else:
|
|
result = response_format.model_validate(json_data)
|
|
else:
|
|
result = content
|
|
|
|
# Extract usage
|
|
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
|
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
|
total_tokens = input_tokens + output_tokens
|
|
|
|
# Record metrics
|
|
duration = time.time() - start_time
|
|
metrics = get_metrics_collector()
|
|
metrics.record_llm_call(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
scope=scope,
|
|
duration=duration,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
success=True,
|
|
)
|
|
|
|
# Record trace span
|
|
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
|
|
|
span_recorder = get_span_recorder()
|
|
span_recorder.record_llm_call(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
scope=scope,
|
|
messages=messages,
|
|
response_content=_serialize_for_span(result),
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
duration=duration,
|
|
finish_reason=finish_reason,
|
|
error=None,
|
|
)
|
|
|
|
if duration > 10.0:
|
|
logger.info(
|
|
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
|
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
|
f"time={duration:.3f}s"
|
|
)
|
|
|
|
if return_usage:
|
|
token_usage = TokenUsage(
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
total_tokens=total_tokens,
|
|
)
|
|
return result, token_usage
|
|
return result
|
|
|
|
except OutputTooLongError:
|
|
raise
|
|
|
|
except json.JSONDecodeError as e:
|
|
last_exception = e
|
|
if attempt < max_retries:
|
|
logger.warning("LiteLLM returned invalid JSON, retrying...")
|
|
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
|
await asyncio.sleep(backoff)
|
|
continue
|
|
else:
|
|
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
|
raise
|
|
|
|
except Exception as e:
|
|
error_str = str(e).lower()
|
|
# Fast fail on auth errors
|
|
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
|
logger.error(f"LiteLLM auth error, not retrying: {e}")
|
|
raise
|
|
|
|
last_exception = e
|
|
if attempt < max_retries:
|
|
# Retry on rate limits, connection errors, server errors
|
|
is_retryable = any(
|
|
keyword in error_str
|
|
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
|
)
|
|
if is_retryable:
|
|
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
|
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
|
await asyncio.sleep(backoff + jitter)
|
|
continue
|
|
|
|
logger.error(f"LiteLLM API error after {attempt + 1} attempts: {e}")
|
|
raise
|
|
|
|
if last_exception:
|
|
raise last_exception
|
|
raise RuntimeError("LiteLLM call failed after all retries")
|
|
|
|
async def call_with_tools(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]],
|
|
max_completion_tokens: int | None = None,
|
|
temperature: float | None = None,
|
|
scope: str = "tools",
|
|
max_retries: int = 5,
|
|
initial_backoff: float = 1.0,
|
|
max_backoff: float = 30.0,
|
|
tool_choice: str | dict[str, Any] = "auto",
|
|
) -> LLMToolCallResult:
|
|
start_time = time.time()
|
|
|
|
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
|
call_kwargs["tools"] = tools
|
|
call_kwargs["tool_choice"] = tool_choice
|
|
|
|
last_exception = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
response = await self._litellm.acompletion(**call_kwargs)
|
|
|
|
message = response.choices[0].message
|
|
content = message.content
|
|
finish_reason = response.choices[0].finish_reason
|
|
|
|
# Extract tool calls
|
|
tool_calls: list[LLMToolCall] = []
|
|
if message.tool_calls:
|
|
for tc in message.tool_calls:
|
|
arguments = tc.function.arguments
|
|
if isinstance(arguments, str):
|
|
arguments = json.loads(arguments)
|
|
tool_calls.append(
|
|
LLMToolCall(
|
|
id=tc.id,
|
|
name=tc.function.name,
|
|
arguments=arguments,
|
|
)
|
|
)
|
|
|
|
# Extract usage
|
|
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
|
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
|
|
|
# Record metrics
|
|
duration = time.time() - start_time
|
|
metrics = get_metrics_collector()
|
|
metrics.record_llm_call(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
scope=scope,
|
|
duration=duration,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
success=True,
|
|
)
|
|
|
|
# Record trace span
|
|
from hindsight_api.tracing import get_span_recorder
|
|
|
|
span_recorder = get_span_recorder()
|
|
tool_calls_dict = (
|
|
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
|
|
if tool_calls
|
|
else None
|
|
)
|
|
span_recorder.record_llm_call(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
scope=scope,
|
|
messages=messages,
|
|
response_content=content,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
duration=duration,
|
|
finish_reason=finish_reason,
|
|
error=None,
|
|
tool_calls=tool_calls_dict,
|
|
)
|
|
|
|
return LLMToolCallResult(
|
|
content=content,
|
|
tool_calls=tool_calls,
|
|
finish_reason=finish_reason or ("tool_calls" if tool_calls else "stop"),
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
)
|
|
|
|
except Exception as e:
|
|
error_str = str(e).lower()
|
|
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
|
raise
|
|
|
|
last_exception = e
|
|
if attempt < max_retries:
|
|
is_retryable = any(
|
|
keyword in error_str
|
|
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
|
)
|
|
if is_retryable:
|
|
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
|
continue
|
|
|
|
logger.error(f"LiteLLM tool call error after {attempt + 1} attempts: {e}")
|
|
raise
|
|
|
|
if last_exception:
|
|
raise last_exception
|
|
raise RuntimeError("LiteLLM tool call failed after all retries")
|
|
|
|
async def cleanup(self) -> None:
|
|
"""Clean up resources."""
|
|
pass
|