2.1 (#2)
Final release of version 2.1. Headlights: ### Core Features - 🤖 Interactive chat with 300+ AI models via OpenRouter - 🔍 Model selection with search and filtering - 💾 Conversation save/load/export (Markdown, JSON, HTML) - 📎 File attachments (images, PDFs, code files) - 💰 Real-time cost tracking and credit monitoring - 🎨 Rich terminal UI with syntax highlighting - 📝 Persistent command history with search (Ctrl+R) - 🌐 Online mode (web search capabilities) - 🧠 Conversation memory toggle ### MCP Integration - 🔧 **File Mode**: AI can read, search, and list local files - Automatic .gitignore filtering - Virtual environment exclusion - Large file handling (auto-truncates >50KB) - ✍️ **Write Mode**: AI can modify files with permission - Create, edit, delete files - Move, copy, organize files - Always requires explicit opt-in - 🗄️ **Database Mode**: AI can query SQLite databases - Read-only access (safe) - Schema inspection - Full SQL query support Reviewed-on: #2 Co-authored-by: Rune Olsen <rune@rune.pm> Co-committed-by: Rune Olsen <rune@rune.pm>
This commit was merged in pull request #2.
This commit is contained in:
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
Core functionality for oAI.
|
||||
|
||||
This module provides the main session management and AI client
|
||||
classes that power the chat application.
|
||||
"""
|
||||
|
||||
from oai.core.session import ChatSession
|
||||
from oai.core.client import AIClient
|
||||
|
||||
__all__ = [
|
||||
"ChatSession",
|
||||
"AIClient",
|
||||
]
|
||||
@@ -0,0 +1,422 @@
|
||||
"""
|
||||
AI Client for oAI.
|
||||
|
||||
This module provides a high-level client for interacting with AI models
|
||||
through the provider abstraction layer.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any, Callable, Dict, Iterator, List, Optional, Union
|
||||
|
||||
from oai.constants import APP_NAME, APP_URL, MODEL_PRICING
|
||||
from oai.providers.base import (
|
||||
AIProvider,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ModelInfo,
|
||||
StreamChunk,
|
||||
ToolCall,
|
||||
UsageStats,
|
||||
)
|
||||
from oai.providers.openrouter import OpenRouterProvider
|
||||
from oai.utils.logging import get_logger
|
||||
|
||||
|
||||
class AIClient:
|
||||
"""
|
||||
High-level AI client for chat interactions.
|
||||
|
||||
Provides a simplified interface for sending chat requests,
|
||||
handling streaming, and managing tool calls.
|
||||
|
||||
Attributes:
|
||||
provider: The underlying AI provider
|
||||
default_model: Default model ID to use
|
||||
http_headers: Custom HTTP headers for requests
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: Optional[str] = None,
|
||||
provider_class: type = OpenRouterProvider,
|
||||
app_name: str = APP_NAME,
|
||||
app_url: str = APP_URL,
|
||||
):
|
||||
"""
|
||||
Initialize the AI client.
|
||||
|
||||
Args:
|
||||
api_key: API key for authentication
|
||||
base_url: Optional custom base URL
|
||||
provider_class: Provider class to use (default: OpenRouterProvider)
|
||||
app_name: Application name for headers
|
||||
app_url: Application URL for headers
|
||||
"""
|
||||
self.provider: AIProvider = provider_class(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
app_name=app_name,
|
||||
app_url=app_url,
|
||||
)
|
||||
self.default_model: Optional[str] = None
|
||||
self.logger = get_logger()
|
||||
|
||||
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
|
||||
"""
|
||||
Get available models.
|
||||
|
||||
Args:
|
||||
filter_text_only: Whether to exclude video-only models
|
||||
|
||||
Returns:
|
||||
List of ModelInfo objects
|
||||
"""
|
||||
return self.provider.list_models(filter_text_only=filter_text_only)
|
||||
|
||||
def get_model(self, model_id: str) -> Optional[ModelInfo]:
|
||||
"""
|
||||
Get information about a specific model.
|
||||
|
||||
Args:
|
||||
model_id: Model identifier
|
||||
|
||||
Returns:
|
||||
ModelInfo or None if not found
|
||||
"""
|
||||
return self.provider.get_model(model_id)
|
||||
|
||||
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get raw model data for provider-specific fields.
|
||||
|
||||
Args:
|
||||
model_id: Model identifier
|
||||
|
||||
Returns:
|
||||
Raw model dictionary or None
|
||||
"""
|
||||
if hasattr(self.provider, "get_raw_model"):
|
||||
return self.provider.get_raw_model(model_id)
|
||||
return None
|
||||
|
||||
def chat(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
model: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
online: bool = False,
|
||||
transforms: Optional[List[str]] = None,
|
||||
) -> Union[ChatResponse, Iterator[StreamChunk]]:
|
||||
"""
|
||||
Send a chat request.
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries
|
||||
model: Model ID (uses default if not specified)
|
||||
stream: Whether to stream the response
|
||||
max_tokens: Maximum tokens in response
|
||||
temperature: Sampling temperature
|
||||
tools: Tool definitions for function calling
|
||||
tool_choice: Tool selection mode
|
||||
system_prompt: System prompt to prepend
|
||||
online: Whether to enable online mode
|
||||
transforms: List of transforms (e.g., ["middle-out"])
|
||||
|
||||
Returns:
|
||||
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
|
||||
|
||||
Raises:
|
||||
ValueError: If no model specified and no default set
|
||||
"""
|
||||
model_id = model or self.default_model
|
||||
if not model_id:
|
||||
raise ValueError("No model specified and no default set")
|
||||
|
||||
# Apply online mode suffix
|
||||
if online and hasattr(self.provider, "get_effective_model_id"):
|
||||
model_id = self.provider.get_effective_model_id(model_id, True)
|
||||
|
||||
# Convert dict messages to ChatMessage objects
|
||||
chat_messages = []
|
||||
|
||||
# Add system prompt if provided
|
||||
if system_prompt:
|
||||
chat_messages.append(ChatMessage(role="system", content=system_prompt))
|
||||
|
||||
# Convert message dicts
|
||||
for msg in messages:
|
||||
# Convert tool_calls dicts to ToolCall objects if present
|
||||
tool_calls_data = msg.get("tool_calls")
|
||||
tool_calls_obj = None
|
||||
if tool_calls_data:
|
||||
from oai.providers.base import ToolCall, ToolFunction
|
||||
tool_calls_obj = []
|
||||
for tc in tool_calls_data:
|
||||
# Handle both ToolCall objects and dicts
|
||||
if isinstance(tc, ToolCall):
|
||||
tool_calls_obj.append(tc)
|
||||
elif isinstance(tc, dict):
|
||||
func_data = tc.get("function", {})
|
||||
tool_calls_obj.append(
|
||||
ToolCall(
|
||||
id=tc.get("id", ""),
|
||||
type=tc.get("type", "function"),
|
||||
function=ToolFunction(
|
||||
name=func_data.get("name", ""),
|
||||
arguments=func_data.get("arguments", "{}"),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
chat_messages.append(
|
||||
ChatMessage(
|
||||
role=msg.get("role", "user"),
|
||||
content=msg.get("content"),
|
||||
tool_calls=tool_calls_obj,
|
||||
tool_call_id=msg.get("tool_call_id"),
|
||||
)
|
||||
)
|
||||
|
||||
self.logger.debug(
|
||||
f"Sending chat request: model={model_id}, "
|
||||
f"messages={len(chat_messages)}, stream={stream}"
|
||||
)
|
||||
|
||||
return self.provider.chat(
|
||||
model=model_id,
|
||||
messages=chat_messages,
|
||||
stream=stream,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
transforms=transforms,
|
||||
)
|
||||
|
||||
def chat_with_tools(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
tools: List[Dict[str, Any]],
|
||||
tool_executor: Callable[[str, Dict[str, Any]], Dict[str, Any]],
|
||||
model: Optional[str] = None,
|
||||
max_loops: int = 5,
|
||||
max_tokens: Optional[int] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
on_tool_call: Optional[Callable[[ToolCall], None]] = None,
|
||||
on_tool_result: Optional[Callable[[str, Dict[str, Any]], None]] = None,
|
||||
) -> ChatResponse:
|
||||
"""
|
||||
Send a chat request with automatic tool call handling.
|
||||
|
||||
Executes tool calls returned by the model and continues
|
||||
the conversation until no more tool calls are requested.
|
||||
|
||||
Args:
|
||||
messages: Initial messages
|
||||
tools: Tool definitions
|
||||
tool_executor: Function to execute tool calls
|
||||
model: Model ID
|
||||
max_loops: Maximum tool call iterations
|
||||
max_tokens: Maximum response tokens
|
||||
system_prompt: System prompt
|
||||
on_tool_call: Callback when tool is called
|
||||
on_tool_result: Callback when tool returns result
|
||||
|
||||
Returns:
|
||||
Final ChatResponse after all tool calls complete
|
||||
"""
|
||||
model_id = model or self.default_model
|
||||
if not model_id:
|
||||
raise ValueError("No model specified and no default set")
|
||||
|
||||
# Build initial messages
|
||||
chat_messages = []
|
||||
if system_prompt:
|
||||
chat_messages.append({"role": "system", "content": system_prompt})
|
||||
chat_messages.extend(messages)
|
||||
|
||||
loop_count = 0
|
||||
current_response: Optional[ChatResponse] = None
|
||||
|
||||
while loop_count < max_loops:
|
||||
# Send request
|
||||
response = self.chat(
|
||||
messages=chat_messages,
|
||||
model=model_id,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
if not isinstance(response, ChatResponse):
|
||||
raise ValueError("Expected non-streaming response")
|
||||
|
||||
current_response = response
|
||||
|
||||
# Check for tool calls
|
||||
tool_calls = response.tool_calls
|
||||
if not tool_calls:
|
||||
break
|
||||
|
||||
self.logger.info(f"Model requested {len(tool_calls)} tool call(s)")
|
||||
|
||||
# Process each tool call
|
||||
tool_results = []
|
||||
for tc in tool_calls:
|
||||
if on_tool_call:
|
||||
on_tool_call(tc)
|
||||
|
||||
try:
|
||||
args = json.loads(tc.function.arguments)
|
||||
except json.JSONDecodeError as e:
|
||||
self.logger.error(f"Failed to parse tool arguments: {e}")
|
||||
result = {"error": f"Invalid arguments: {e}"}
|
||||
else:
|
||||
result = tool_executor(tc.function.name, args)
|
||||
|
||||
if on_tool_result:
|
||||
on_tool_result(tc.function.name, result)
|
||||
|
||||
tool_results.append({
|
||||
"tool_call_id": tc.id,
|
||||
"role": "tool",
|
||||
"name": tc.function.name,
|
||||
"content": json.dumps(result),
|
||||
})
|
||||
|
||||
# Add assistant message with tool calls
|
||||
assistant_msg = {
|
||||
"role": "assistant",
|
||||
"content": response.content,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": tc.type,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in tool_calls
|
||||
],
|
||||
}
|
||||
chat_messages.append(assistant_msg)
|
||||
chat_messages.extend(tool_results)
|
||||
|
||||
loop_count += 1
|
||||
|
||||
if loop_count >= max_loops:
|
||||
self.logger.warning(f"Reached max tool call loops ({max_loops})")
|
||||
|
||||
return current_response
|
||||
|
||||
def stream_chat(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
model: Optional[str] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
online: bool = False,
|
||||
on_chunk: Optional[Callable[[StreamChunk], None]] = None,
|
||||
) -> tuple[str, Optional[UsageStats]]:
|
||||
"""
|
||||
Stream a chat response and collect the full text.
|
||||
|
||||
Args:
|
||||
messages: Chat messages
|
||||
model: Model ID
|
||||
max_tokens: Maximum tokens
|
||||
system_prompt: System prompt
|
||||
online: Online mode
|
||||
on_chunk: Optional callback for each chunk
|
||||
|
||||
Returns:
|
||||
Tuple of (full_response_text, usage_stats)
|
||||
"""
|
||||
response = self.chat(
|
||||
messages=messages,
|
||||
model=model,
|
||||
stream=True,
|
||||
max_tokens=max_tokens,
|
||||
system_prompt=system_prompt,
|
||||
online=online,
|
||||
)
|
||||
|
||||
if isinstance(response, ChatResponse):
|
||||
# Not actually streaming
|
||||
return response.content or "", response.usage
|
||||
|
||||
full_text = ""
|
||||
usage: Optional[UsageStats] = None
|
||||
|
||||
for chunk in response:
|
||||
if chunk.error:
|
||||
self.logger.error(f"Stream error: {chunk.error}")
|
||||
break
|
||||
|
||||
if chunk.delta_content:
|
||||
full_text += chunk.delta_content
|
||||
if on_chunk:
|
||||
on_chunk(chunk)
|
||||
|
||||
if chunk.usage:
|
||||
usage = chunk.usage
|
||||
|
||||
return full_text, usage
|
||||
|
||||
def get_credits(self) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get account credit information.
|
||||
|
||||
Returns:
|
||||
Credit info dict or None if unavailable
|
||||
"""
|
||||
return self.provider.get_credits()
|
||||
|
||||
def estimate_cost(
|
||||
self,
|
||||
model_id: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
) -> float:
|
||||
"""
|
||||
Estimate cost for a completion.
|
||||
|
||||
Args:
|
||||
model_id: Model ID
|
||||
input_tokens: Number of input tokens
|
||||
output_tokens: Number of output tokens
|
||||
|
||||
Returns:
|
||||
Estimated cost in USD
|
||||
"""
|
||||
if hasattr(self.provider, "estimate_cost"):
|
||||
return self.provider.estimate_cost(model_id, input_tokens, output_tokens)
|
||||
|
||||
# Fallback to default pricing
|
||||
input_cost = MODEL_PRICING["input"] * input_tokens / 1_000_000
|
||||
output_cost = MODEL_PRICING["output"] * output_tokens / 1_000_000
|
||||
return input_cost + output_cost
|
||||
|
||||
def set_default_model(self, model_id: str) -> None:
|
||||
"""
|
||||
Set the default model.
|
||||
|
||||
Args:
|
||||
model_id: Model ID to use as default
|
||||
"""
|
||||
self.default_model = model_id
|
||||
self.logger.info(f"Default model set to: {model_id}")
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""Clear the provider's model cache."""
|
||||
if hasattr(self.provider, "clear_cache"):
|
||||
self.provider.clear_cache()
|
||||
@@ -0,0 +1,659 @@
|
||||
"""
|
||||
Chat session management for oAI.
|
||||
|
||||
This module provides the ChatSession class that manages an interactive
|
||||
chat session including history, state, and message handling.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
from rich.live import Live
|
||||
from rich.markdown import Markdown
|
||||
|
||||
from oai.commands.registry import CommandContext, CommandResult, registry
|
||||
from oai.config.database import Database
|
||||
from oai.config.settings import Settings
|
||||
from oai.constants import (
|
||||
COST_WARNING_THRESHOLD,
|
||||
LOW_CREDIT_AMOUNT,
|
||||
LOW_CREDIT_RATIO,
|
||||
)
|
||||
from oai.core.client import AIClient
|
||||
from oai.mcp.manager import MCPManager
|
||||
from oai.providers.base import ChatResponse, StreamChunk, UsageStats
|
||||
from oai.ui.console import (
|
||||
console,
|
||||
display_markdown,
|
||||
display_panel,
|
||||
print_error,
|
||||
print_info,
|
||||
print_success,
|
||||
print_warning,
|
||||
)
|
||||
from oai.ui.prompts import prompt_copy_response
|
||||
from oai.utils.logging import get_logger
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionStats:
|
||||
"""
|
||||
Statistics for the current session.
|
||||
|
||||
Tracks tokens, costs, and message counts.
|
||||
"""
|
||||
|
||||
total_input_tokens: int = 0
|
||||
total_output_tokens: int = 0
|
||||
total_cost: float = 0.0
|
||||
message_count: int = 0
|
||||
|
||||
@property
|
||||
def total_tokens(self) -> int:
|
||||
"""Get total token count."""
|
||||
return self.total_input_tokens + self.total_output_tokens
|
||||
|
||||
def add_usage(self, usage: Optional[UsageStats], cost: float = 0.0) -> None:
|
||||
"""
|
||||
Add usage stats from a response.
|
||||
|
||||
Args:
|
||||
usage: Usage statistics
|
||||
cost: Cost if not in usage
|
||||
"""
|
||||
if usage:
|
||||
self.total_input_tokens += usage.prompt_tokens
|
||||
self.total_output_tokens += usage.completion_tokens
|
||||
if usage.total_cost_usd:
|
||||
self.total_cost += usage.total_cost_usd
|
||||
else:
|
||||
self.total_cost += cost
|
||||
else:
|
||||
self.total_cost += cost
|
||||
self.message_count += 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class HistoryEntry:
|
||||
"""
|
||||
A single entry in the conversation history.
|
||||
|
||||
Stores the user prompt, assistant response, and metrics.
|
||||
"""
|
||||
|
||||
prompt: str
|
||||
response: str
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
msg_cost: float = 0.0
|
||||
timestamp: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary format."""
|
||||
return {
|
||||
"prompt": self.prompt,
|
||||
"response": self.response,
|
||||
"prompt_tokens": self.prompt_tokens,
|
||||
"completion_tokens": self.completion_tokens,
|
||||
"msg_cost": self.msg_cost,
|
||||
}
|
||||
|
||||
|
||||
class ChatSession:
|
||||
"""
|
||||
Manages an interactive chat session.
|
||||
|
||||
Handles conversation history, state management, command processing,
|
||||
and communication with the AI client.
|
||||
|
||||
Attributes:
|
||||
client: AI client for API requests
|
||||
settings: Application settings
|
||||
mcp_manager: MCP manager for file/database access
|
||||
history: Conversation history
|
||||
stats: Session statistics
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: AIClient,
|
||||
settings: Settings,
|
||||
mcp_manager: Optional[MCPManager] = None,
|
||||
):
|
||||
"""
|
||||
Initialize a chat session.
|
||||
|
||||
Args:
|
||||
client: AI client instance
|
||||
settings: Application settings
|
||||
mcp_manager: Optional MCP manager
|
||||
"""
|
||||
self.client = client
|
||||
self.settings = settings
|
||||
self.mcp_manager = mcp_manager
|
||||
self.db = Database()
|
||||
|
||||
self.history: List[HistoryEntry] = []
|
||||
self.stats = SessionStats()
|
||||
|
||||
# Session state
|
||||
self.system_prompt: str = settings.effective_system_prompt
|
||||
self.memory_enabled: bool = True
|
||||
self.memory_start_index: int = 0
|
||||
self.online_enabled: bool = settings.default_online_mode
|
||||
self.middle_out_enabled: bool = False
|
||||
self.session_max_token: int = 0
|
||||
self.current_index: int = 0
|
||||
|
||||
# Selected model
|
||||
self.selected_model: Optional[Dict[str, Any]] = None
|
||||
|
||||
self.logger = get_logger()
|
||||
|
||||
def get_context(self) -> CommandContext:
|
||||
"""
|
||||
Get the current command context.
|
||||
|
||||
Returns:
|
||||
CommandContext with current session state
|
||||
"""
|
||||
return CommandContext(
|
||||
settings=self.settings,
|
||||
provider=self.client.provider,
|
||||
mcp_manager=self.mcp_manager,
|
||||
selected_model_raw=self.selected_model,
|
||||
session_history=[e.to_dict() for e in self.history],
|
||||
session_system_prompt=self.system_prompt,
|
||||
memory_enabled=self.memory_enabled,
|
||||
memory_start_index=self.memory_start_index,
|
||||
online_enabled=self.online_enabled,
|
||||
middle_out_enabled=self.middle_out_enabled,
|
||||
session_max_token=self.session_max_token,
|
||||
total_input_tokens=self.stats.total_input_tokens,
|
||||
total_output_tokens=self.stats.total_output_tokens,
|
||||
total_cost=self.stats.total_cost,
|
||||
message_count=self.stats.message_count,
|
||||
current_index=self.current_index,
|
||||
)
|
||||
|
||||
def set_model(self, model: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Set the selected model.
|
||||
|
||||
Args:
|
||||
model: Raw model dictionary
|
||||
"""
|
||||
self.selected_model = model
|
||||
self.client.set_default_model(model["id"])
|
||||
self.logger.info(f"Model selected: {model['id']}")
|
||||
|
||||
def build_api_messages(self, user_input: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Build the messages array for an API request.
|
||||
|
||||
Includes system prompt, history (if memory enabled), and current input.
|
||||
|
||||
Args:
|
||||
user_input: Current user input
|
||||
|
||||
Returns:
|
||||
List of message dictionaries
|
||||
"""
|
||||
messages = []
|
||||
|
||||
# Add system prompt
|
||||
if self.system_prompt:
|
||||
messages.append({"role": "system", "content": self.system_prompt})
|
||||
|
||||
# Add database context if in database mode
|
||||
if self.mcp_manager and self.mcp_manager.enabled:
|
||||
if self.mcp_manager.mode == "database" and self.mcp_manager.selected_db_index is not None:
|
||||
db = self.mcp_manager.databases[self.mcp_manager.selected_db_index]
|
||||
db_context = (
|
||||
f"You are connected to SQLite database: {db['name']}\n"
|
||||
f"Available tables: {', '.join(db['tables'])}\n\n"
|
||||
"Use inspect_database, search_database, or query_database tools. "
|
||||
"All queries are read-only."
|
||||
)
|
||||
messages.append({"role": "system", "content": db_context})
|
||||
|
||||
# Add history if memory enabled
|
||||
if self.memory_enabled:
|
||||
for i in range(self.memory_start_index, len(self.history)):
|
||||
entry = self.history[i]
|
||||
messages.append({"role": "user", "content": entry.prompt})
|
||||
messages.append({"role": "assistant", "content": entry.response})
|
||||
|
||||
# Add current message
|
||||
messages.append({"role": "user", "content": user_input})
|
||||
|
||||
return messages
|
||||
|
||||
def get_mcp_tools(self) -> Optional[List[Dict[str, Any]]]:
|
||||
"""
|
||||
Get MCP tool definitions if available.
|
||||
|
||||
Returns:
|
||||
List of tool schemas or None
|
||||
"""
|
||||
if not self.mcp_manager or not self.mcp_manager.enabled:
|
||||
return None
|
||||
|
||||
if not self.selected_model:
|
||||
return None
|
||||
|
||||
# Check if model supports tools
|
||||
supported_params = self.selected_model.get("supported_parameters", [])
|
||||
if "tools" not in supported_params and "functions" not in supported_params:
|
||||
return None
|
||||
|
||||
return self.mcp_manager.get_tools_schema()
|
||||
|
||||
async def execute_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
tool_args: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute an MCP tool.
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool
|
||||
tool_args: Tool arguments
|
||||
|
||||
Returns:
|
||||
Tool execution result
|
||||
"""
|
||||
if not self.mcp_manager:
|
||||
return {"error": "MCP not available"}
|
||||
|
||||
return await self.mcp_manager.call_tool(tool_name, **tool_args)
|
||||
|
||||
def send_message(
|
||||
self,
|
||||
user_input: str,
|
||||
stream: bool = True,
|
||||
on_stream_chunk: Optional[Callable[[str], None]] = None,
|
||||
) -> Tuple[str, Optional[UsageStats], float]:
|
||||
"""
|
||||
Send a message and get a response.
|
||||
|
||||
Args:
|
||||
user_input: User's input text
|
||||
stream: Whether to stream the response
|
||||
on_stream_chunk: Callback for stream chunks
|
||||
|
||||
Returns:
|
||||
Tuple of (response_text, usage_stats, response_time)
|
||||
"""
|
||||
if not self.selected_model:
|
||||
raise ValueError("No model selected")
|
||||
|
||||
start_time = time.time()
|
||||
messages = self.build_api_messages(user_input)
|
||||
|
||||
# Get MCP tools
|
||||
tools = self.get_mcp_tools()
|
||||
if tools:
|
||||
# Disable streaming when tools are present
|
||||
stream = False
|
||||
|
||||
# Build request parameters
|
||||
model_id = self.selected_model["id"]
|
||||
if self.online_enabled:
|
||||
if hasattr(self.client.provider, "get_effective_model_id"):
|
||||
model_id = self.client.provider.get_effective_model_id(model_id, True)
|
||||
|
||||
transforms = ["middle-out"] if self.middle_out_enabled else None
|
||||
|
||||
max_tokens = None
|
||||
if self.session_max_token > 0:
|
||||
max_tokens = self.session_max_token
|
||||
|
||||
if tools:
|
||||
# Use tool handling flow
|
||||
response = self._send_with_tools(
|
||||
messages=messages,
|
||||
model_id=model_id,
|
||||
tools=tools,
|
||||
max_tokens=max_tokens,
|
||||
transforms=transforms,
|
||||
)
|
||||
response_time = time.time() - start_time
|
||||
return response.content or "", response.usage, response_time
|
||||
|
||||
elif stream:
|
||||
# Use streaming flow
|
||||
full_text, usage = self._stream_response(
|
||||
messages=messages,
|
||||
model_id=model_id,
|
||||
max_tokens=max_tokens,
|
||||
transforms=transforms,
|
||||
on_chunk=on_stream_chunk,
|
||||
)
|
||||
response_time = time.time() - start_time
|
||||
return full_text, usage, response_time
|
||||
|
||||
else:
|
||||
# Non-streaming request
|
||||
response = self.client.chat(
|
||||
messages=messages,
|
||||
model=model_id,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
transforms=transforms,
|
||||
)
|
||||
response_time = time.time() - start_time
|
||||
|
||||
if isinstance(response, ChatResponse):
|
||||
return response.content or "", response.usage, response_time
|
||||
else:
|
||||
return "", None, response_time
|
||||
|
||||
def _send_with_tools(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
model_id: str,
|
||||
tools: List[Dict[str, Any]],
|
||||
max_tokens: Optional[int] = None,
|
||||
transforms: Optional[List[str]] = None,
|
||||
) -> ChatResponse:
|
||||
"""
|
||||
Send a request with tool call handling.
|
||||
|
||||
Args:
|
||||
messages: API messages
|
||||
model_id: Model ID
|
||||
tools: Tool definitions
|
||||
max_tokens: Max tokens
|
||||
transforms: Transforms list
|
||||
|
||||
Returns:
|
||||
Final ChatResponse
|
||||
"""
|
||||
max_loops = 5
|
||||
loop_count = 0
|
||||
api_messages = list(messages)
|
||||
|
||||
while loop_count < max_loops:
|
||||
response = self.client.chat(
|
||||
messages=api_messages,
|
||||
model=model_id,
|
||||
stream=False,
|
||||
max_tokens=max_tokens,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
transforms=transforms,
|
||||
)
|
||||
|
||||
if not isinstance(response, ChatResponse):
|
||||
raise ValueError("Expected ChatResponse")
|
||||
|
||||
tool_calls = response.tool_calls
|
||||
if not tool_calls:
|
||||
return response
|
||||
|
||||
console.print(f"\n[dim yellow]🔧 AI requesting {len(tool_calls)} tool call(s)...[/]")
|
||||
|
||||
tool_results = []
|
||||
for tc in tool_calls:
|
||||
try:
|
||||
args = json.loads(tc.function.arguments)
|
||||
except json.JSONDecodeError as e:
|
||||
self.logger.error(f"Failed to parse tool arguments: {e}")
|
||||
tool_results.append({
|
||||
"tool_call_id": tc.id,
|
||||
"role": "tool",
|
||||
"name": tc.function.name,
|
||||
"content": json.dumps({"error": f"Invalid arguments: {e}"}),
|
||||
})
|
||||
continue
|
||||
|
||||
# Display tool call
|
||||
args_display = ", ".join(
|
||||
f'{k}="{v}"' if isinstance(v, str) else f"{k}={v}"
|
||||
for k, v in args.items()
|
||||
)
|
||||
console.print(f"[dim cyan] → {tc.function.name}({args_display})[/]")
|
||||
|
||||
# Execute tool
|
||||
result = asyncio.run(self.execute_tool(tc.function.name, args))
|
||||
|
||||
if "error" in result:
|
||||
console.print(f"[dim red] ✗ Error: {result['error']}[/]")
|
||||
else:
|
||||
self._display_tool_success(tc.function.name, result)
|
||||
|
||||
tool_results.append({
|
||||
"tool_call_id": tc.id,
|
||||
"role": "tool",
|
||||
"name": tc.function.name,
|
||||
"content": json.dumps(result),
|
||||
})
|
||||
|
||||
# Add assistant message with tool calls
|
||||
api_messages.append({
|
||||
"role": "assistant",
|
||||
"content": response.content,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": tc.type,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in tool_calls
|
||||
],
|
||||
})
|
||||
api_messages.extend(tool_results)
|
||||
|
||||
console.print("\n[dim cyan]💭 Processing tool results...[/]")
|
||||
loop_count += 1
|
||||
|
||||
self.logger.warning(f"Reached max tool loops ({max_loops})")
|
||||
console.print(f"[bold yellow]⚠️ Reached maximum tool calls ({max_loops})[/]")
|
||||
return response
|
||||
|
||||
def _display_tool_success(self, tool_name: str, result: Dict[str, Any]) -> None:
|
||||
"""Display a success message for a tool call."""
|
||||
if tool_name == "search_files":
|
||||
count = result.get("count", 0)
|
||||
console.print(f"[dim green] ✓ Found {count} file(s)[/]")
|
||||
elif tool_name == "read_file":
|
||||
size = result.get("size", 0)
|
||||
truncated = " (truncated)" if result.get("truncated") else ""
|
||||
console.print(f"[dim green] ✓ Read {size} bytes{truncated}[/]")
|
||||
elif tool_name == "list_directory":
|
||||
count = result.get("count", 0)
|
||||
console.print(f"[dim green] ✓ Listed {count} item(s)[/]")
|
||||
elif tool_name == "inspect_database":
|
||||
if "table" in result:
|
||||
console.print(f"[dim green] ✓ Inspected table: {result['table']}[/]")
|
||||
else:
|
||||
console.print(f"[dim green] ✓ Inspected database ({result.get('table_count', 0)} tables)[/]")
|
||||
elif tool_name == "search_database":
|
||||
count = result.get("count", 0)
|
||||
console.print(f"[dim green] ✓ Found {count} match(es)[/]")
|
||||
elif tool_name == "query_database":
|
||||
count = result.get("count", 0)
|
||||
console.print(f"[dim green] ✓ Query returned {count} row(s)[/]")
|
||||
else:
|
||||
console.print("[dim green] ✓ Success[/]")
|
||||
|
||||
def _stream_response(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
model_id: str,
|
||||
max_tokens: Optional[int] = None,
|
||||
transforms: Optional[List[str]] = None,
|
||||
on_chunk: Optional[Callable[[str], None]] = None,
|
||||
) -> Tuple[str, Optional[UsageStats]]:
|
||||
"""
|
||||
Stream a response with live display.
|
||||
|
||||
Args:
|
||||
messages: API messages
|
||||
model_id: Model ID
|
||||
max_tokens: Max tokens
|
||||
transforms: Transforms
|
||||
on_chunk: Callback for chunks
|
||||
|
||||
Returns:
|
||||
Tuple of (full_text, usage)
|
||||
"""
|
||||
response = self.client.chat(
|
||||
messages=messages,
|
||||
model=model_id,
|
||||
stream=True,
|
||||
max_tokens=max_tokens,
|
||||
transforms=transforms,
|
||||
)
|
||||
|
||||
if isinstance(response, ChatResponse):
|
||||
return response.content or "", response.usage
|
||||
|
||||
full_text = ""
|
||||
usage: Optional[UsageStats] = None
|
||||
|
||||
try:
|
||||
with Live("", console=console, refresh_per_second=10) as live:
|
||||
for chunk in response:
|
||||
if chunk.error:
|
||||
console.print(f"\n[bold red]Stream error: {chunk.error}[/]")
|
||||
break
|
||||
|
||||
if chunk.delta_content:
|
||||
full_text += chunk.delta_content
|
||||
live.update(Markdown(full_text))
|
||||
if on_chunk:
|
||||
on_chunk(chunk.delta_content)
|
||||
|
||||
if chunk.usage:
|
||||
usage = chunk.usage
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[bold yellow]⚠️ Streaming interrupted[/]")
|
||||
return "", None
|
||||
|
||||
return full_text, usage
|
||||
|
||||
def add_to_history(
|
||||
self,
|
||||
prompt: str,
|
||||
response: str,
|
||||
usage: Optional[UsageStats] = None,
|
||||
cost: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Add an exchange to the history.
|
||||
|
||||
Args:
|
||||
prompt: User prompt
|
||||
response: Assistant response
|
||||
usage: Usage statistics
|
||||
cost: Cost if not in usage
|
||||
"""
|
||||
entry = HistoryEntry(
|
||||
prompt=prompt,
|
||||
response=response,
|
||||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||||
completion_tokens=usage.completion_tokens if usage else 0,
|
||||
msg_cost=usage.total_cost_usd if usage and usage.total_cost_usd else cost,
|
||||
timestamp=time.time(),
|
||||
)
|
||||
self.history.append(entry)
|
||||
self.current_index = len(self.history) - 1
|
||||
self.stats.add_usage(usage, cost)
|
||||
|
||||
def save_conversation(self, name: str) -> bool:
|
||||
"""
|
||||
Save the current conversation.
|
||||
|
||||
Args:
|
||||
name: Name for the saved conversation
|
||||
|
||||
Returns:
|
||||
True if saved successfully
|
||||
"""
|
||||
if not self.history:
|
||||
return False
|
||||
|
||||
data = [e.to_dict() for e in self.history]
|
||||
self.db.save_conversation(name, data)
|
||||
self.logger.info(f"Saved conversation: {name}")
|
||||
return True
|
||||
|
||||
def load_conversation(self, name: str) -> bool:
|
||||
"""
|
||||
Load a saved conversation.
|
||||
|
||||
Args:
|
||||
name: Name of the conversation to load
|
||||
|
||||
Returns:
|
||||
True if loaded successfully
|
||||
"""
|
||||
data = self.db.load_conversation(name)
|
||||
if not data:
|
||||
return False
|
||||
|
||||
self.history.clear()
|
||||
for entry_dict in data:
|
||||
self.history.append(HistoryEntry(
|
||||
prompt=entry_dict.get("prompt", ""),
|
||||
response=entry_dict.get("response", ""),
|
||||
prompt_tokens=entry_dict.get("prompt_tokens", 0),
|
||||
completion_tokens=entry_dict.get("completion_tokens", 0),
|
||||
msg_cost=entry_dict.get("msg_cost", 0.0),
|
||||
))
|
||||
|
||||
self.current_index = len(self.history) - 1
|
||||
self.memory_start_index = 0
|
||||
self.stats = SessionStats() # Reset stats for loaded conversation
|
||||
self.logger.info(f"Loaded conversation: {name}")
|
||||
return True
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the session state."""
|
||||
self.history.clear()
|
||||
self.stats = SessionStats()
|
||||
self.system_prompt = ""
|
||||
self.memory_start_index = 0
|
||||
self.current_index = 0
|
||||
self.logger.info("Session reset")
|
||||
|
||||
def check_warnings(self) -> List[str]:
|
||||
"""
|
||||
Check for cost and credit warnings.
|
||||
|
||||
Returns:
|
||||
List of warning messages
|
||||
"""
|
||||
warnings = []
|
||||
|
||||
# Check last message cost
|
||||
if self.history:
|
||||
last_cost = self.history[-1].msg_cost
|
||||
threshold = self.settings.cost_warning_threshold
|
||||
if last_cost > threshold:
|
||||
warnings.append(
|
||||
f"High cost: ${last_cost:.4f} exceeds threshold ${threshold:.4f}"
|
||||
)
|
||||
|
||||
# Check credits
|
||||
credits = self.client.get_credits()
|
||||
if credits:
|
||||
left = credits.get("credits_left", 0)
|
||||
total = credits.get("total_credits", 0)
|
||||
|
||||
if left < LOW_CREDIT_AMOUNT:
|
||||
warnings.append(f"Low credits: ${left:.2f} remaining!")
|
||||
elif total > 0 and left < total * LOW_CREDIT_RATIO:
|
||||
warnings.append(f"Credits low: less than 10% remaining (${left:.2f})")
|
||||
|
||||
return warnings
|
||||
Reference in New Issue
Block a user