"""Chat service module""" import json import uuid import logging from typing import List, Dict, Any, AsyncGenerator, Optional from luxx.models import Conversation, Message from luxx.tools.executor import ToolExecutor from luxx.tools.core import registry from luxx.services.llm_client import LLMClient from luxx.config import config logger = logging.getLogger(__name__) # Maximum iterations to prevent infinite loops MAX_ITERATIONS = 10 def _sse_event(event: str, data: dict) -> str: """Format a Server-Sent Event string.""" return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n" def get_llm_client(conversation: Conversation = None): """Get LLM client, optionally using conversation's provider. Returns (client, max_tokens)""" max_tokens = None if conversation and conversation.provider_id: from luxx.models import LLMProvider from luxx.database import SessionLocal db = SessionLocal() try: provider = db.query(LLMProvider).filter(LLMProvider.id == conversation.provider_id).first() if provider: max_tokens = provider.max_tokens client = LLMClient( api_key=provider.api_key, api_url=provider.base_url, model=provider.default_model ) return client, max_tokens finally: db.close() # Fallback to global config client = LLMClient() return client, max_tokens class StreamContext: """Context for streaming response state management.""" def __init__( self, step_index: int = 0, current_step_id: str = None, current_step_idx: int = None, current_stream_type: str = None, full_content: str = "", full_thinking: str = "" ): self.step_index = step_index self.current_step_id = current_step_id self.current_step_idx = current_step_idx self.current_stream_type = current_stream_type self.full_content = full_content self.full_thinking = full_thinking self.all_steps = [] self.all_tool_calls = [] self.all_tool_results = [] self.tool_calls_list = [] def reset_iteration(self): """Reset streaming step tracker for new iteration.""" self.current_step_id = None self.current_step_idx = None self.current_stream_type = None self.full_content = "" self.full_thinking = "" self.tool_calls_list = [] def start_stream_step(self, step_type: str) -> str: """Start a new streaming step. Returns the step_id.""" self.current_step_idx = self.step_index self.current_step_id = f"step-{self.step_index}" self.current_stream_type = step_type self.step_index += 1 return self.current_step_id def yield_stream_step(self, step_type: str, content: str) -> Dict[str, Any]: """Yield a streaming step event.""" return _sse_event("process_step", { "step": { "id": self.current_step_id, "index": self.current_step_idx, "type": step_type, "content": content } }) def save_streaming_step(self): """Save the current streaming step to all_steps.""" if self.current_step_id is None: return if self.current_stream_type == "thinking": self.all_steps.append({ "id": self.current_step_id, "index": self.current_step_idx, "type": "thinking", "content": self.full_thinking }) elif self.current_stream_type == "text": self.all_steps.append({ "id": self.current_step_id, "index": self.current_step_idx, "type": "text", "content": self.full_content }) def handle_thinking_stream(self, delta: Dict) -> Optional[Dict]: """Handle reasoning/thinking delta. Returns yield_obj if step was yielded.""" reasoning = delta.get("reasoning_content", "") if not reasoning: return None prev_len = len(self.full_thinking) self.full_thinking += reasoning if prev_len == 0: # New thinking stream started self.start_stream_step("thinking") return self.yield_stream_step("thinking", self.full_thinking) def handle_text_stream(self, delta: Dict) -> Optional[Dict]: """Handle content delta. Returns yield_obj if step was yielded.""" content = delta.get("content", "") if not content: return None prev_len = len(self.full_content) self.full_content += content if prev_len == 0: # New text stream started self.start_stream_step("text") return self.yield_stream_step("text", self.full_content) def handle_tool_call(self) -> tuple: """Handle tool calls. Returns (tool_call_step_ids, tool_call_steps, yield_objs).""" tool_call_step_ids = [] tool_call_steps = [] yield_objs = [] for tc in self.tool_calls_list: call_step_idx = self.step_index call_step_id = f"step-{self.step_index}" tool_call_step_ids.append(call_step_id) self.step_index += 1 call_step = { "id": call_step_id, "index": call_step_idx, "type": "tool_call", "id_ref": tc.get("id", ""), "name": tc["function"]["name"], "arguments": tc["function"]["arguments"] } tool_call_steps.append(call_step) yield_objs.append(_sse_event("process_step", {"step": call_step})) return tool_call_step_ids, tool_call_steps, yield_objs def handle_tool_result(self, tool_result: Dict, tool_call_step_id: str) -> tuple: """Handle single tool result. Returns (result_step, yield_obj).""" result_step_idx = self.step_index result_step_id = f"step-{self.step_index}" self.step_index += 1 content = tool_result.get("content", "") success = True try: content_obj = json.loads(content) if isinstance(content_obj, dict): success = content_obj.get("success", True) except: pass result_step = { "id": result_step_id, "index": result_step_idx, "type": "tool_result", "id_ref": tool_call_step_id, "name": tool_result.get("name", ""), "content": content, "success": success } return result_step, _sse_event("process_step", {"step": result_step}) class ChatService: """Chat service with tool support""" def __init__(self): self.tool_executor = ToolExecutor() def build_messages( self, conversation: Conversation, include_system: bool = True ) -> List[Dict[str, str]]: """Build message list""" from luxx.database import SessionLocal from luxx.models import Message messages = [] if include_system and conversation.system_prompt: messages.append({ "role": "system", "content": conversation.system_prompt }) db = SessionLocal() try: db_messages = db.query(Message).filter( Message.conversation_id == conversation.id ).order_by(Message.created_at).all() for msg in db_messages: # Parse JSON content if possible try: content_obj = json.loads(msg.content) if msg.content else {} if isinstance(content_obj, dict): content = content_obj.get("text", msg.content) else: content = msg.content except (json.JSONDecodeError, TypeError): content = msg.content messages.append({ "role": msg.role, "content": content }) finally: db.close() return messages async def stream_response( self, conversation: Conversation, user_message: str, thinking_enabled: bool = False, enabled_tools: list = None, user_id: int = None, username: str = None, workspace: str = None, user_permission_level: int = 1 ) -> AsyncGenerator[Dict[str, str], None]: """ Streaming response generator Yields raw SSE event strings for direct forwarding. """ try: messages = self.build_messages(conversation) messages.append({ "role": "user", "content": json.dumps({"text": user_message, "attachments": []}) }) # Get tools based on enabled_tools filter if enabled_tools: tools = [t for t in registry.list_all() if t.get("function", {}).get("name") in enabled_tools] else: tools = [] llm, provider_max_tokens = get_llm_client(conversation) model = conversation.model or llm.default_model or "gpt-4" # 直接使用 provider 的 max_tokens max_tokens = provider_max_tokens # Token usage tracking total_usage = { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 } actual_token_count = 0 # Streaming context for state management ctx = StreamContext() for iteration in range(MAX_ITERATIONS): # Reset streaming context for this iteration ctx.reset_iteration() async for sse_line in llm.stream_call( model=model, messages=messages, tools=tools, temperature=conversation.temperature, max_tokens=max_tokens or 8192, thinking_enabled=thinking_enabled or conversation.thinking_enabled ): # Parse SSE line # Format: "event: xxx\ndata: {...}\n\n" event_type = None data_str = None for line in sse_line.strip().split('\n'): if line.startswith('event: '): event_type = line[7:].strip() elif line.startswith('data: '): data_str = line[6:].strip() if data_str is None: continue # Handle error events from LLM if event_type == 'error': try: error_data = json.loads(data_str) yield _sse_event("error", {"content": error_data.get("content", "Unknown error")}) except json.JSONDecodeError: yield _sse_event("error", {"content": data_str}) return # Parse the data try: chunk = json.loads(data_str) except json.JSONDecodeError: yield _sse_event("error", {"content": f"Failed to parse response: {data_str}"}) return # 提取 API 返回的 usage 信息 if "usage" in chunk: usage = chunk["usage"] total_usage["prompt_tokens"] = usage.get("prompt_tokens", 0) total_usage["completion_tokens"] = usage.get("completion_tokens", 0) total_usage["total_tokens"] = usage.get("total_tokens", 0) # Check for error in response if "error" in chunk: error_msg = chunk["error"].get("message", str(chunk["error"])) yield _sse_event("error", {"content": f"API Error: {error_msg}"}) return # Get delta choices = chunk.get("choices", []) if not choices: # Check if there's any content in the response (for non-standard LLM responses) if chunk.get("content") or chunk.get("message"): content = chunk.get("content") or chunk.get("message", {}).get("content", "") if content: prev_len = len(ctx.full_content) ctx.full_content += content if prev_len == 0: # New text stream started ctx.start_stream_step("text") yield _sse_event("process_step", { "step": { "id": ctx.current_step_id if prev_len == 0 else f"step-{ctx.step_index - 1}", "index": ctx.current_step_idx if prev_len == 0 else ctx.step_index - 1, "type": "text", "content": ctx.full_content } }) continue delta = choices[0].get("delta", {}) # Handle reasoning (thinking) yield_obj = ctx.handle_thinking_stream(delta) if yield_obj: yield yield_obj # Handle content yield_obj = ctx.handle_text_stream(delta) if yield_obj: yield yield_obj # Accumulate tool calls tool_calls_delta = delta.get("tool_calls", []) for tc in tool_calls_delta: idx = tc.get("index", 0) if idx >= len(ctx.tool_calls_list): ctx.tool_calls_list.append({ "id": tc.get("id", ""), "type": "function", "function": {"name": "", "arguments": ""} }) func = tc.get("function", {}) if func.get("name"): ctx.tool_calls_list[idx]["function"]["name"] += func["name"] if func.get("arguments"): ctx.tool_calls_list[idx]["function"]["arguments"] += func["arguments"] # Save streaming step (thinking or text) ctx.save_streaming_step() # Handle tool calls if ctx.tool_calls_list: ctx.all_tool_calls.extend(ctx.tool_calls_list) # Handle tool_call steps tool_call_step_ids, tool_call_steps, yield_objs = ctx.handle_tool_call() ctx.all_steps.extend(tool_call_steps) for yield_obj in yield_objs: yield yield_obj # Execute tools tool_context = { "workspace": workspace, "user_id": user_id, "username": username, "user_permission_level": user_permission_level } tool_results = self.tool_executor.process_tool_calls_parallel( ctx.tool_calls_list, tool_context ) # Handle tool_result steps for i, tr in enumerate(tool_results): tool_call_step_id = tool_call_step_ids[i] if i < len(tool_call_step_ids) else f"step-{i}" result_step, yield_obj = ctx.handle_tool_result(tr, tool_call_step_id) ctx.all_steps.append(result_step) yield yield_obj ctx.all_tool_results.append({ "role": "tool", "tool_call_id": tr.get("tool_call_id", ""), "content": tr.get("content", "") }) # Add assistant message with tool calls for next iteration messages.append({ "role": "assistant", "content": ctx.full_content or "", "tool_calls": ctx.tool_calls_list }) messages.extend(ctx.all_tool_results[-len(tool_results):]) ctx.all_tool_results = [] continue # No tool calls - final iteration, save message msg_id = str(uuid.uuid4()) # 使用 API 返回的真实 completion_tokens，如果 API 没返回则降级使用估算值 actual_token_count = total_usage.get("completion_tokens", 0) or len(ctx.full_content) // 4 logger.info(f"[TOKEN] total_usage: {total_usage}, actual_token_count: {actual_token_count}") self._save_message( conversation.id, msg_id, ctx.full_content, ctx.all_tool_calls, ctx.all_tool_results, ctx.all_steps, actual_token_count, total_usage ) yield _sse_event("done", { "message_id": msg_id, "token_count": actual_token_count, "usage": total_usage }) return # Max iterations exceeded - save message before error if ctx.full_content or ctx.all_tool_calls: msg_id = str(uuid.uuid4()) self._save_message( conversation.id, msg_id, ctx.full_content, ctx.all_tool_calls, ctx.all_tool_results, ctx.all_steps, actual_token_count, total_usage ) yield _sse_event("error", {"content": "Exceeded maximum tool call iterations"}) except Exception as e: logger.error(f"Stream error: {e}") yield _sse_event("error", {"content": str(e)}) def _save_message( self, conversation_id: str, msg_id: str, full_content: str, all_tool_calls: list, all_tool_results: list, all_steps: list, token_count: int = 0, usage: dict = None ): """Save the assistant message to database.""" from luxx.database import SessionLocal from luxx.models import Message content_json = { "text": full_content, "steps": all_steps } if all_tool_calls: content_json["tool_calls"] = all_tool_calls db = SessionLocal() try: msg = Message( id=msg_id, conversation_id=conversation_id, role="assistant", content=json.dumps(content_json, ensure_ascii=False), token_count=token_count, usage=json.dumps(usage) if usage else None ) db.add(msg) db.commit() except Exception as e: db.rollback() raise finally: db.close() # Global chat service chat_service = ChatService()