codemate-cli 1.0.0__py3-none-any.whl

codemate/client.py

@@ -0,0 +1,311 @@
+import httpx
+import json
+from typing import AsyncGenerator, Dict, List, Optional, Any, Tuple
+from pathlib import Path
+
+from codemate.config import Config
+
+
+class CodeMateClient:
+    """Client for interacting with the CodeMate API"""
+    
+    def __init__(self, config: Config):
+        self.config = config
+        self.base_url = config.get_endpoint()
+        self.api_key = config.get_api_key()
+        self.timeout = httpx.Timeout(120.0, connect=10.0)
+        self._kb_cache = None  # Cache for knowledge bases
+    
+    def _get_headers(self) -> Dict[str, str]:
+        """Get request headers with authorization"""
+        return {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}",
+        }
+    
+    def _prepare_messages(
+        self,
+        message: str,
+        conversation_history: Optional[List[Dict]] = None,
+        context: Optional[List[Dict]] = None
+    ) -> List[Dict]:
+        """Prepare messages array for API request"""
+        messages = conversation_history.copy() if conversation_history else []
+        
+        # Add current message with context if provided
+        user_message = {
+            "role": "user",
+            "content": message
+        }
+        
+        # If context is provided, add it to the message along with web_search flag
+        if context:
+            user_message["context"] = context
+            user_message["web_search"] = False
+        else:
+            # Even without context, add empty context and web_search false
+            user_message["context"] = []
+            user_message["web_search"] = False
+        
+        messages.append(user_message)
+        
+        return messages
+    
+    async def chat_stream(
+        self,
+        message: str,
+        model: str = "chat_c0_cli",
+        conversation_history: Optional[List[Dict]] = None,
+        conversation_id: Optional[str] = None,
+        call_for: str = "chat_c0_cli",
+        tools: Optional[List[Dict]] = None,
+        context: Optional[List[Dict]] = None,
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        """
+        Stream chat completions from the API
+        
+        Yields dictionaries with:
+        - type: "message" | "reasoning" | "tool_call" | "tool_result" | "metadata"
+        - content: the actual content
+        - Additional fields depending on type
+        """
+        messages = self._prepare_messages(message, conversation_history, context)
+        
+        payload = {
+            "messages": messages,
+            "call_for": "chat_c0_cli",
+            "mode": "NORMAL",
+            "conversation_id": conversation_id,
+        }
+        
+        client = httpx.AsyncClient(timeout=self.timeout)
+        stream_context = None
+        reasoning_content = ""
+        message_content = ""
+        
+        try:
+            stream_context = client.stream(
+                "POST",
+                f"{self.base_url}/chat/stream",
+                headers=self._get_headers(),
+                json=payload,
+            )
+            response = await stream_context.__aenter__()
+            response.raise_for_status()
+
+            buffer = ""
+            should_stop = False
+            
+            async for chunk in response.aiter_text():
+                # with open("extra.json", "a", encoding="utf-8") as f:
+                #     f.write(chunk)
+                #     f.write("\n")
+
+                if should_stop:
+                    break
+                    
+                buffer += chunk
+                
+                while "<__!!__END__!!__>" in buffer:
+                    msg_str, buffer = buffer.split("<__!!__END__!!__>", 1)
+                    msg_str = msg_str.strip()
+                    
+                    if not msg_str:
+                        continue
+
+                    try:
+                        chunk_data = json.loads(msg_str)
+                        msg_type = chunk_data.get("type")
+
+                        # ADD ERROR HANDLING:
+                        if msg_type == "error":
+                            yield {
+                                "type": "error",
+                                "message": chunk_data.get("message", "Unknown error")
+                            }
+                            should_stop = True
+                            break
+                        
+                        elif msg_type == "end":
+                            should_stop = True
+                            break
+                        elif msg_type == "message":
+                            content = chunk_data.get("message")
+                            reasoning = chunk_data.get("reasoning_content")
+                            if content:
+                                message_content += content
+                                yield {"type": "message", "content": content}
+                            elif reasoning:
+                                reasoning_content += reasoning
+                                yield {"type": "reasoning", "content": reasoning}
+                        elif msg_type == "tool_calls":
+                            # Yield tool calls information
+                            yield {
+                                "type": "tool_calls",
+                                "tool_calls": chunk_data.get("tool_calls", [])
+                            }
+                        elif msg_type == "tool_call_start":
+                            # Yield tool call start
+                            yield {
+                                "type": "tool_call_start",
+                                "tool_call": chunk_data.get("tool_call", {})
+                            }
+                        elif msg_type == "tool_call_complete":
+                            # Yield tool call completion with results
+                            yield {
+                                "type": "tool_call_complete",
+                                "tool_call": chunk_data.get("tool_call", {})
+                            }
+                        elif msg_type == "usage":
+                            # Yield usage information
+                            yield {
+                                "type": "usage",
+                                "usage": chunk_data.get("usage", {})
+                            }
+
+                    except json.JSONDecodeError as e:
+                        print(f"JSON decode error: {e}, chunk: {msg_str[:100]}")
+                        continue
+            
+            # Process remaining buffer
+            if buffer.strip() and not should_stop:
+                try:
+                    chunk_data = json.loads(buffer.strip())
+                    msg_type = chunk_data.get("type")
+                    if msg_type == "message":
+                        content = chunk_data.get("message")
+                        reasoning = chunk_data.get("reasoning_content")
+                        if content:
+                            message_content += content
+                        elif reasoning:
+                            reasoning_content += reasoning
+                except json.JSONDecodeError:
+                    pass
+
+            # Write to files after stream completes
+            # with open("reasoning_content.txt", "w", encoding="utf-8") as f:
+            #     f.write(reasoning_content)
+            # with open("message_content.txt", "w", encoding="utf-8") as f:
+            #     f.write(message_content)
+                    
+        finally:
+            # ✅ Ensure cleanup happens
+            if stream_context:
+                try:
+                    await stream_context.__aexit__(None, None, None)
+                except Exception:
+                    pass
+            
+            try:
+                await client.aclose()
+            except Exception:
+                pass
+
+
+    def chat(
+        self,
+        message: str,
+        model: str = "chat_c0_cli",
+        conversation_history: Optional[List[Dict]] = None,
+        conversation_id: Optional[str] = None,
+        call_for: str = "chat_c0_cli",
+        tools: Optional[List[Dict]] = None,
+        context: Optional[List[Dict]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Send a chat completion request (non-streaming)
+        
+        Returns the complete response
+        """
+        messages = self._prepare_messages(message, conversation_history, context)
+        
+        payload = {
+            "messages": messages,
+            "call_for": "chat_c0_cli",
+            "mode": "NORMAL",
+            "conversation_id": conversation_id,
+        }
+        
+        with httpx.Client(timeout=self.timeout) as client:
+            response = client.post(
+                f"{self.base_url}/chat/stream",
+                headers=self._get_headers(),
+                json=payload,
+            )
+            response.raise_for_status()
+            return response.json()
+    
+    async def get_models(self) -> List[Dict[str, Any]]:
+        """Get list of available models"""
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.get(
+                f"{self.base_url}/v1/models",
+                headers=self._get_headers(),
+            )
+            response.raise_for_status()
+            return response.json()
+
+    async def list_kbs(self) -> List[Dict[str, Any]]:
+        """Get list of knowledge bases"""
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.get(
+                f"{self.base_url}/list_kbs?include_cloud=false",
+                headers=self._get_headers(),
+            )
+            response.raise_for_status()
+            data = response.json()
+            # Cache the knowledge bases
+            self._kb_cache = data
+            return data
+    
+    def get_cached_kbs(self) -> Optional[List[Dict[str, Any]]]:
+        """Get cached knowledge bases"""
+        return self._kb_cache
+
+
+# Synchronous wrapper for CLI usage
+class SyncCodeMateClient:
+    """Synchronous wrapper for the async client"""
+    
+    def __init__(self, config: Config):
+        self.async_client = CodeMateClient(config)
+    
+    def chat_stream(self, *args, **kwargs):
+        """Stream chat - returns async generator"""
+        return self.async_client.chat_stream(*args, **kwargs)
+    
+    def chat(self, *args, **kwargs):
+        """Non-streaming chat"""
+        return self.async_client.chat(*args, **kwargs)
+
+    def list_kbs(self):
+        """List knowledge bases"""
+        import asyncio
+        
+        try:
+            # Check if there's already a running event loop
+            loop = asyncio.get_running_loop()
+            # If we're in an async context, we need to create a new task and wait for it
+            import concurrent.futures
+            import threading
+            
+            def run_in_thread():
+                """Run the async method in a separate thread with its own event loop"""
+                new_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(new_loop)
+                try:
+                    return new_loop.run_until_complete(self.async_client.list_kbs())
+                finally:
+                    new_loop.close()
+            
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(run_in_thread)
+                return future.result()
+                
+        except RuntimeError:
+            # No running event loop, safe to use asyncio.run()
+            return asyncio.run(self.async_client.list_kbs())
+    
+    def get_cached_kbs(self):
+        """Get cached knowledge bases"""
+        return self.async_client.get_cached_kbs()

codemate/commands/__init__.py

@@ -0,0 +1,6 @@
+"""
+Command handlers for CodeMate CLI
+"""
+
+# This module will contain individual command implementations
+# Commands are defined in cli.py using Click decorators

codemate/commands/config.py

@@ -0,0 +1,103 @@
+"""
+Configuration command implementation
+"""
+
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+
+from codemate.config import Config
+
+
+class ConfigCommand:
+    """Handler for configuration commands"""
+    
+    def __init__(self, config: Config, console: Console):
+        self.config = config
+        self.console = console
+    
+    def set_api_key(self, api_key: str):
+        """Set API key"""
+        if not api_key or not api_key.strip():
+            self.console.print("[red]Error: API key cannot be empty[/red]")
+            return False
+        
+        self.config.set_api_key(api_key)
+        self.console.print("[green]✓ API key saved successfully![/green]")
+        return True
+    
+    def get_api_key(self):
+        """Display masked API key"""
+        key = self.config.get_api_key()
+        if key:
+            masked = key[:8] + "..." + key[-4:] if len(key) > 12 else "***"
+            self.console.print(f"API Key: [cyan]{masked}[/cyan]")
+        else:
+            self.console.print("[yellow]No API key configured.[/yellow]")
+            self.console.print(
+                "Set one with: [cyan]codemate config set-key YOUR_API_KEY[/cyan]"
+            )
+    
+    def set_endpoint(self, endpoint: str):
+        """Set API endpoint"""
+        if not endpoint or not endpoint.strip():
+            self.console.print("[red]Error: Endpoint cannot be empty[/red]")
+            return False
+        
+        # Basic URL validation
+        if not endpoint.startswith(("http://", "https://")):
+            self.console.print(
+                "[yellow]Warning: Endpoint should start with http:// or https://[/yellow]"
+            )
+        
+        self.config.set_endpoint(endpoint)
+        self.console.print(f"[green]✓ Endpoint set to: {endpoint}[/green]")
+        return True
+    
+    def get_endpoint(self):
+        """Display current endpoint"""
+        endpoint = self.config.get_endpoint()
+        self.console.print(f"Endpoint: [cyan]{endpoint}[/cyan]")
+    
+    def set_model(self, model: str):
+        """Set default model"""
+        self.config.set_default_model(model)
+        self.console.print(f"[green]✓ Default model set to: {model}[/green]")
+    
+    def show_all(self):
+        """Display all configuration settings"""
+        settings = self.config.get_all()
+        
+        table = Table(
+            title="Current Configuration",
+            show_header=True,
+            header_style="bold magenta"
+        )
+        table.add_column("Setting", style="cyan", width=20)
+        table.add_column("Value", style="white")
+        
+        for key, value in settings.items():
+            # Mask API key
+            if key == "api_key" and value:
+                display_value = value[:8] + "..." + value[-4:] if len(value) > 12 else "***"
+            else:
+                display_value = str(value)
+            
+            table.add_row(key, display_value)
+        
+        self.console.print(table)
+    
+    def reset(self):
+        """Reset configuration to defaults"""
+        self.console.print("[yellow]Resetting configuration to defaults...[/yellow]")
+        self.config.reset()
+        self.console.print("[green]✓ Configuration reset successfully![/green]")
+    
+    def show_location(self):
+        """Show configuration file location"""
+        self.console.print(Panel(
+            f"Configuration file location:\n[cyan]{self.config.config_file}[/cyan]\n\n"
+            f"History directory:\n[cyan]{self.config.config_dir / 'history'}[/cyan]",
+            title="[bold]Config Location[/bold]",
+            border_style="cyan"
+        ))