iridet-bot 0.1.1a1__py3-none-any.whl

iribot/main.py

@@ -0,0 +1,300 @@
+"""Main FastAPI application"""
+import json
+from pathlib import Path
+from fastapi import FastAPI, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+
+from .config import settings
+from .models import (
+    SessionCreate,
+    SystemPromptUpdate,
+    ChatRequest,
+    MessageRecord,
+    ToolCallRecord,
+    SystemPromptGenerateRequest,
+    SystemPromptGenerateResponse,
+)
+from .session_manager import session_manager
+from .agent import agent
+from .executor import tool_executor
+from .prompt_generator import generate_system_prompt
+
+# Initialize FastAPI app
+app = FastAPI(
+    title=settings.app_title,
+    debug=settings.debug
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.cors_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# Prompt Generation Endpoints
+@app.post("/api/prompt/generate")
+def generate_prompt(request: SystemPromptGenerateRequest):
+    """
+    Generate a system prompt with current date/time, safety policies, and tool descriptions
+    
+    Args:
+        request: GeneratePromptRequest with optional custom_instructions
+        
+    Returns:
+        Generated system prompt with metadata
+    """
+    from .prompt_generator import get_current_datetime_info
+    
+    prompt = generate_system_prompt(
+        custom_instructions=request.custom_instructions
+    )
+    
+    return SystemPromptGenerateResponse(
+        system_prompt=prompt,
+        datetime_info=get_current_datetime_info()
+    )
+
+
+@app.get("/api/prompt/generate")
+def generate_prompt_get(custom_instructions: str = ""):
+    """
+    Generate a system prompt (GET endpoint for convenience)
+    
+    Query Parameters:
+        custom_instructions: Optional custom instructions
+        
+    Returns:
+        Generated system prompt with metadata
+    """
+    from .prompt_generator import get_current_datetime_info
+    
+    prompt = generate_system_prompt(
+        custom_instructions=custom_instructions
+    )
+    
+    return SystemPromptGenerateResponse(
+        system_prompt=prompt,
+        datetime_info=get_current_datetime_info()
+    )
+
+
+@app.post("/api/sessions")
+def create_session(request: SessionCreate):
+    """Create a new chat session"""
+    # If system_prompt not provided, generate one automatically
+    system_prompt = request.system_prompt
+    if system_prompt is None:
+        system_prompt = generate_system_prompt(custom_instructions="")
+    
+    session = session_manager.create_session(
+        title=request.title,
+        system_prompt=system_prompt
+    )
+    return session.model_dump()
+
+
+@app.get("/api/sessions")
+def list_sessions():
+    """List all chat sessions"""
+    sessions = session_manager.list_sessions()
+    return [s.model_dump() for s in sessions]
+
+
+@app.get("/api/sessions/{session_id}")
+def get_session(session_id: str):
+    """Get a specific session with all records"""
+    session = session_manager.get_session(session_id)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    
+    return session.model_dump()
+
+
+@app.delete("/api/sessions/{session_id}")
+def delete_session(session_id: str):
+    """Delete a session"""
+    success = session_manager.delete_session(session_id)
+    if not success:
+        raise HTTPException(status_code=404, detail="Session not found")
+    return {"status": "success"}
+
+
+@app.get("/api/tools/status")
+def get_tools_status():
+    """Get status of all tools"""
+    return tool_executor.get_all_tool_statuses()
+
+
+# Streaming chat endpoint
+@app.post("/api/chat/stream")
+def chat_stream(request: ChatRequest):
+    """Send a message and get AI response with streaming support"""
+    
+    def generate():
+        # Get session
+        session = session_manager.get_session(request.session_id)
+        if not session:
+            yield f"data: {json.dumps({'type': 'error', 'content': 'Session not found'})}\n\n"
+            return
+        
+        # Add user message record
+        user_record = MessageRecord(
+            role="user",
+            content=request.message,
+            binary_content=request.binary_content
+        )
+        session_manager.add_record(request.session_id, user_record.model_dump())
+        
+        # Send user record event
+        yield f"data: {json.dumps({'type': 'record', 'record': user_record.model_dump()}, default=str)}\n\n"
+        
+        # Get messages for LLM
+        messages = session_manager.get_messages_for_llm(request.session_id)
+        
+        try:
+            # Tool calling loop
+            max_iterations = 50
+            
+            for iteration in range(max_iterations):
+                # Stream AI response
+                current_content = ""
+                tool_calls = []
+                
+                for chunk in agent.chat_stream(
+                    messages=messages[1:],
+                    system_prompt=session.system_prompt,
+                    images=[bc.get("data") for bc in (request.binary_content or []) if bc.get("data")] if iteration == 0 else None
+                ):
+                    if chunk["type"] == "content":
+                        current_content += chunk["content"]
+                        yield f"data: {json.dumps({'type': 'content', 'content': chunk['content']})}\n\n"
+                    elif chunk["type"] == "done":
+                        tool_calls = chunk.get("tool_calls", [])
+                
+                if not tool_calls:
+                    # No tool calls, save and finish
+                    if current_content:
+                        assistant_record = MessageRecord(
+                            role="assistant",
+                            content=current_content
+                        )
+                        session_manager.add_record(request.session_id, assistant_record.model_dump())
+                        yield f"data: {json.dumps({'type': 'record', 'record': assistant_record.model_dump()}, default=str)}\n\n"
+                    
+                    yield f"data: {json.dumps({'type': 'done'})}\n\n"
+                    return
+                
+                # Has tool calls - save thinking message if any
+                if current_content:
+                    thinking_record = MessageRecord(
+                        role="assistant",
+                        content=current_content
+                    )
+                    session_manager.add_record(request.session_id, thinking_record.model_dump())
+                    yield f"data: {json.dumps({'type': 'record', 'record': thinking_record.model_dump()}, default=str)}\n\n"
+                    messages.append({"role": "assistant", "content": current_content, "tool_calls": tool_calls})
+                else:
+                    messages.append({"role": "assistant", "content": "", "tool_calls": tool_calls})
+                
+                # Signal tool calls starting
+                yield f"data: {json.dumps({'type': 'tool_calls_start', 'tool_calls': tool_calls})}\n\n"
+                
+                # Execute tool calls
+                for tool_call in tool_calls:
+                    tool_name = tool_call["function"]["name"]
+                    tool_args_str = tool_call["function"]["arguments"]
+                    tool_id = tool_call["id"]
+                    
+                    # Parse arguments
+                    try:
+                        tool_args = json.loads(tool_args_str) if isinstance(tool_args_str, str) else tool_args_str
+                    except json.JSONDecodeError:
+                        tool_args = {"raw_arguments": tool_args_str}
+                    
+                    # Signal tool execution starting
+                    yield f"data: {json.dumps({'type': 'tool_start', 'tool_call_id': tool_id, 'tool_name': tool_name, 'arguments': tool_args, 'success': None})}\n\n"
+                    
+                    # Execute the tool
+                    result = agent.process_tool_call(
+                        tool_name,
+                        tool_args_str,
+                        context={"session_id": request.session_id}
+                    )
+                    
+                    # Create and save tool call record
+                    tool_record = ToolCallRecord(
+                        tool_call_id=tool_id,
+                        tool_name=tool_name,
+                        arguments=tool_args,
+                        result=result.get("result") if result.get("success") else result.get("error"),
+                        success=result.get("success", False)
+                    )
+                    session_manager.add_record(request.session_id, tool_record.model_dump())
+                    
+                    # Send tool result
+                    yield f"data: {json.dumps({'type': 'tool_result', 'record': tool_record.model_dump()}, default=str)}\n\n"
+                    
+                    # Add tool result to messages
+                    messages.append({
+                        "role": "tool",
+                        "tool_call_id": tool_id,
+                        "content": json.dumps(result) if isinstance(result, dict) else str(result)
+                    })
+            
+            # Max iterations reached
+            final_content = "Tool execution reached maximum iterations. Please try again with a simpler request."
+            final_record = MessageRecord(role="assistant", content=final_content)
+            session_manager.add_record(request.session_id, final_record.model_dump())
+            yield f"data: {json.dumps({'type': 'record', 'record': final_record.model_dump()}, default=str)}\n\n"
+            yield f"data: {json.dumps({'type': 'done'})}\n\n"
+            
+        except Exception as e:
+            error_content = f"Error: {str(e)}"
+            error_record = MessageRecord(role="assistant", content=error_content)
+            session_manager.add_record(request.session_id, error_record.model_dump())
+            yield f"data: {json.dumps({'type': 'error', 'content': error_content})}\n\n"
+    
+    return StreamingResponse(
+        generate(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no"
+        }
+    )
+
+
+# System prompt endpoints
+@app.post("/api/sessions/{session_id}/system-prompt")
+def update_system_prompt(session_id: str, request: SystemPromptUpdate):
+    """Update session system prompt"""
+    session = session_manager.update_system_prompt(session_id, request.system_prompt)
+    if not session:
+        raise HTTPException(status_code=404, detail="Session not found")
+    return session.model_dump()
+
+
+# Health check
+@app.get("/api/health")
+def health_check():
+    """Health check endpoint"""
+    return {"status": "ok"}
+
+
+# Serve frontend static files
+static_dir = Path(__file__).parent / "static"
+static_dir.mkdir(parents=True, exist_ok=True)
+print(static_dir)
+app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8009)

iribot/models.py

@@ -0,0 +1,79 @@
+"""Data models for the Agent application"""
+from typing import Optional, List, Dict, Any, Literal, Union
+from datetime import datetime
+from pydantic import BaseModel, Field
+
+
+def get_local_now():
+    """Get the current local time."""
+    return datetime.now()
+
+
+# ============ Session Record Types ============
+
+class MessageRecord(BaseModel):
+    """Message record - represents a message in the conversation"""
+    type: Literal["message"] = "message"
+    role: Literal["system", "user", "assistant"]
+    content: str
+    binary_content: Optional[List[Dict[str, Any]]] = None  # Images, files, etc.
+    timestamp: datetime = Field(default_factory=get_local_now)
+
+
+class ToolCallRecord(BaseModel):
+    """Tool call record - represents a single tool invocation and result"""
+    type: Literal["tool_call"] = "tool_call"
+    tool_call_id: str
+    tool_name: str
+    arguments: Dict[str, Any]
+    result: Any
+    success: bool
+    timestamp: datetime = Field(default_factory=get_local_now)
+
+
+# Union type for session records
+SessionRecord = Union[MessageRecord, ToolCallRecord]
+
+
+# ============ Session Model ============
+
+class Session(BaseModel):
+    """Session model with unified record list"""
+    id: str = Field(default_factory=lambda: str(datetime.now().timestamp()))
+    title: str
+    records: List[Dict[str, Any]] = []  # List of MessageRecord or ToolCallRecord
+    created_at: datetime = Field(default_factory=get_local_now)
+    updated_at: datetime = Field(default_factory=get_local_now)
+    system_prompt: str  # Required field, must be provided when creating session
+
+
+# ============ API Request/Response Models ============
+
+class ChatRequest(BaseModel):
+    """Chat request model"""
+    session_id: str
+    message: str
+    binary_content: Optional[List[Dict[str, Any]]] = None
+
+class SystemPromptUpdate(BaseModel):
+    """System prompt update model"""
+    system_prompt: str
+
+
+# ============ Prompt Generation Models ============
+
+class SystemPromptGenerateRequest(BaseModel):
+    """Request to generate a system prompt"""
+    custom_instructions: str = ""
+
+
+class SystemPromptGenerateResponse(BaseModel):
+    """Response containing generated system prompt"""
+    system_prompt: str
+    datetime_info: dict
+
+
+class SessionCreate(BaseModel):
+    """Create new session request"""
+    title: str = "New Chat"
+    system_prompt: Optional[str] = None

iribot/prompt_generator.py

@@ -0,0 +1,104 @@
+"""System Prompt Generator for Agent"""
+
+from datetime import datetime
+from zoneinfo import ZoneInfo
+from typing import List, Dict, Any
+from pathlib import Path
+from jinja2 import Environment, FileSystemLoader
+from .executor import tool_executor
+
+
+# Initialize Jinja2 Environment
+TEMPLATE_DIR = Path(__file__).parent / "templates"
+jinja_env = Environment(
+    loader=FileSystemLoader(str(TEMPLATE_DIR)),
+    trim_blocks=True,
+    lstrip_blocks=True,
+    autoescape=False
+)
+
+
+def get_current_datetime_info() -> Dict[str, str]:
+    """
+    Get current date, time and timezone information
+    
+    Returns:
+        Dictionary containing formatted date/time/timezone
+    """
+    # Get current UTC time
+    utc_now = datetime.now(ZoneInfo("UTC"))
+    
+    # Also get local timezone (system timezone)
+    local_tz = datetime.now().astimezone().tzinfo
+    local_now = datetime.now(local_tz)
+    
+    return {
+        "current_utc": utc_now.strftime("%Y-%m-%d %H:%M:%S UTC"),
+        "current_local": local_now.strftime("%Y-%m-%d %H:%M:%S %Z"),
+        "timezone": str(local_tz),
+    }
+
+
+def get_available_tools_description() -> str:
+    """
+    Get descriptions of all available tools for the prompt
+    
+    Returns:
+        Formatted string describing all available tools
+    """
+    tools = tool_executor.get_all_tools()
+    
+    if not tools:
+        return "No tools are currently available."
+    
+    tools_description = "## Available Tools\n\n"
+    
+    for tool in tools:
+        # Extract tool info from OpenAI function format
+        func = tool.get('function', {})
+        tool_name = func.get('name', 'Unknown')
+        tool_desc = func.get('description', 'No description available')
+        params = func.get('parameters', {})
+        
+        tools_description += f"### {tool_name}\n"
+        tools_description += f"Description: {tool_desc}\n"
+        
+        if 'properties' in params:
+            tools_description += "Parameters:\n"
+            for param_name, param_info in params['properties'].items():
+                param_type = param_info.get('type', 'unknown')
+                param_desc = param_info.get('description', 'No description')
+                required = param_name in params.get('required', [])
+                req_mark = " (required)" if required else " (optional)"
+                tools_description += f"  - `{param_name}` ({param_type}){req_mark}: {param_desc}\n"
+        
+        tools_description += "\n"
+    
+    return tools_description
+
+
+def generate_system_prompt(custom_instructions: str = "") -> str:
+    """
+    Generate a system prompt for the Agent using Jinja2 template
+    
+    Args:
+        custom_instructions: Optional custom instructions to append to the prompt
+    
+    Returns:
+        Complete system prompt string
+    """
+    datetime_info = get_current_datetime_info()
+    tools_desc = get_available_tools_description()
+    
+    # Load and render template
+    template = jinja_env.get_template("system_prompt.j2")
+    
+    prompt = template.render(
+        current_utc=datetime_info['current_utc'],
+        current_local=datetime_info['current_local'],
+        timezone=datetime_info['timezone'],
+        tools_description=tools_desc,
+        custom_instructions=custom_instructions
+    )
+    
+    return prompt

iribot/session_manager.py

@@ -0,0 +1,194 @@
+"""Session management for storing and retrieving chat sessions"""
+from typing import Dict, Optional, List, Any
+from datetime import datetime
+import json
+from pathlib import Path
+from .models import Session, MessageRecord
+
+
+class SessionManager:
+    """Manages chat sessions and persistence"""
+    
+    def __init__(self, storage_path: str = "./sessions"):
+        self.storage_path = Path(storage_path)
+        self.storage_path.mkdir(exist_ok=True)
+        self.sessions: Dict[str, Session] = {}
+        self._load_all_sessions()
+    
+    def create_session(self, title: str = "New Chat", system_prompt: Optional[str] = None) -> Session:
+        """Create a new session with system message"""
+        if system_prompt is None:
+            raise ValueError("system_prompt is required. Please generate a system prompt using the prompt generator first.")
+        
+        session = Session(
+            title=title,
+            system_prompt=system_prompt
+        )
+        # Add system message as first record
+        system_record = MessageRecord(
+            role="system",
+            content=session.system_prompt
+        )
+        session.records.append(system_record.model_dump())
+        
+        self.sessions[session.id] = session
+        self._save_session(session)
+        return session
+    
+    def get_session(self, session_id: str) -> Optional[Session]:
+        """Get a session by ID"""
+        return self.sessions.get(session_id)
+    
+    def list_sessions(self) -> List[Session]:
+        """List all sessions sorted by updated_at descending"""
+        return sorted(
+            self.sessions.values(),
+            key=lambda s: s.updated_at,
+            reverse=True
+        )
+    
+    def add_record(self, session_id: str, record: Dict[str, Any]) -> Optional[Session]:
+        """Add a record to a session"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return None
+        
+        session.records.append(record)
+        session.updated_at = datetime.now()
+        self._save_session(session)
+        return session
+    
+    def add_records(self, session_id: str, records: List[Dict[str, Any]]) -> Optional[Session]:
+        """Add multiple records to a session"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return None
+        
+        session.records.extend(records)
+        session.updated_at = datetime.now()
+        self._save_session(session)
+        return session
+    
+    def delete_session(self, session_id: str) -> bool:
+        """Delete a session"""
+        if session_id not in self.sessions:
+            return False
+        
+        del self.sessions[session_id]
+        session_file = self.storage_path / f"{session_id}.json"
+        if session_file.exists():
+            session_file.unlink()
+        return True
+    
+    def update_system_prompt(self, session_id: str, system_prompt: str) -> Optional[Session]:
+        """Update session system prompt"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return None
+        
+        session.system_prompt = system_prompt
+        # Update the system message record if it exists
+        if session.records and session.records[0].get("type") == "message" and session.records[0].get("role") == "system":
+            session.records[0]["content"] = system_prompt
+        
+        session.updated_at = datetime.now()
+        self._save_session(session)
+        return session
+    
+    def get_messages_for_llm(self, session_id: str) -> List[Dict[str, Any]]:
+        """Extract messages in LLM-compatible format from session records"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return []
+        
+        messages = []
+        
+        for record in session.records:
+            if record.get("type") == "message":
+                role = record.get("role")
+                content = record.get("content", "")
+                
+                if role in ["system", "user"]:
+                    messages.append({"role": role, "content": content})
+                elif role == "assistant":
+                    # Check if this assistant message had tool calls following it
+                    messages.append({"role": "assistant", "content": content})
+            
+            elif record.get("type") == "tool_call":
+                # Add tool call to assistant message and tool result
+                tool_call_id = record.get("tool_call_id")
+                tool_name = record.get("tool_name")
+                arguments = record.get("arguments", {})
+                result = record.get("result")
+                
+                # If last message is assistant, add tool_calls to it
+                if messages and messages[-1]["role"] == "assistant":
+                    if "tool_calls" not in messages[-1]:
+                        messages[-1]["tool_calls"] = []
+                    messages[-1]["tool_calls"].append({
+                        "id": tool_call_id,
+                        "type": "function",
+                        "function": {
+                            "name": tool_name,
+                            "arguments": json.dumps(arguments) if isinstance(arguments, dict) else str(arguments)
+                        }
+                    })
+                
+                # Add tool result message
+                messages.append({
+                    "role": "tool",
+                    "tool_call_id": tool_call_id,
+                    "content": json.dumps(result) if isinstance(result, dict) else str(result)
+                })
+        
+        return messages
+    
+    def _save_session(self, session: Session) -> None:
+        """Save session to disk"""
+        session_file = self.storage_path / f"{session.id}.json"
+        with open(session_file, 'w', encoding='utf-8') as f:
+            json.dump(session.model_dump(), f, indent=2, default=str, ensure_ascii=False)
+    
+    def _load_all_sessions(self) -> None:
+        """Load all sessions from disk"""
+        for session_file in self.storage_path.glob("*.json"):
+            try:
+                with open(session_file, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    # Handle migration from old format
+                    if "messages" in data and "records" not in data:
+                        data = self._migrate_old_format(data)
+                    session = Session(**data)
+                    self.sessions[session.id] = session
+            except Exception as e:
+                print(f"Error loading session {session_file}: {e}")
+    
+    def _migrate_old_format(self, data: Dict) -> Dict:
+        """Migrate from old message format to new record format"""
+        records = []
+        
+        # Add system message
+        records.append({
+            "type": "message",
+            "role": "system",
+            "content": data.get("system_prompt", ""),
+            "timestamp": data.get("created_at")
+        })
+        
+        # Convert old messages to records
+        for msg in data.get("messages", []):
+            records.append({
+                "type": "message",
+                "role": msg.get("role"),
+                "content": msg.get("content", ""),
+                "binary_content": msg.get("binary_content"),
+                "timestamp": msg.get("timestamp")
+            })
+        
+        data["records"] = records
+        del data["messages"]
+        return data
+
+
+# Global session manager instance
+session_manager = SessionManager()