jarvis-ai-assistant 0.1.115__py3-none-any.whl → 0.1.116__py3-none-any.whl

jarvis/jarvis_platform_manager/main.py

@@ -1,5 +1,17 @@
 from jarvis.jarvis_platform.registry import PlatformRegistry
-from jarvis.utils import PrettyOutput, OutputType, init_env, get_multiline_input
+from jarvis.jarvis_utils import PrettyOutput, OutputType, init_env, get_multiline_input
+import asyncio
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+from typing import List, Dict, Any, Optional
+import uvicorn
+import io
+from contextlib import redirect_stdout
+import json
+import os
+from datetime import datetime
+from fastapi.middleware.cors import CORSMiddleware
 
 def list_platforms():
     """List all supported platforms and models"""
@@ -93,17 +105,412 @@ def chat_with_model(platform_name: str, model_name: str):
         except:
             pass
 
-def info_command(args):
-    """Process info subcommand"""
-    list_platforms()
+# Helper function for platform and model validation
+def validate_platform_model(args):
+    if not args.platform or not args.model:
+        PrettyOutput.print("请指定平台和模型。使用 'jarvis info' 查看可用平台和模型。", OutputType.ERROR)
+        return False
+    return True
 
 def chat_command(args):
     """Process chat subcommand"""
-    if not args.platform or not args.model:
-        PrettyOutput.print("请指定平台和模型。使用 'jarvis info' 查看可用平台和模型。", OutputType.ERROR)
+    if not validate_platform_model(args):
         return
     chat_with_model(args.platform, args.model)
 
+def info_command(args):
+    """Process info subcommand"""
+    list_platforms()
+
+# New models for OpenAI-compatible API
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    stream: bool = False
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: ChatMessage
+    finish_reason: str = "stop"
+
+class ChatCompletionChunk(BaseModel):
+    id: str
+    object: str = "chat.completion.chunk"
+    created: int
+    model: str
+    choices: List[Dict[str, Any]]
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Dict[str, int] = Field(default_factory=lambda: {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+
+def service_command(args):
+    """Process service subcommand - start OpenAI-compatible API server"""
+    import time
+    import uuid
+    import json
+    import os
+    from datetime import datetime
+    
+    host = args.host
+    port = args.port
+    default_platform = args.platform
+    default_model = args.model
+    
+    # Create logs directory if it doesn't exist
+    logs_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs")
+    os.makedirs(logs_dir, exist_ok=True)
+    
+    app = FastAPI(title="Jarvis API Server")
+    
+    # 添加 CORS 中间件
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],  # 允许所有来源，生产环境应更严格
+        allow_credentials=True,
+        allow_methods=["*"],  # 允许所有方法
+        allow_headers=["*"],  # 允许所有头
+    )
+    
+    registry = PlatformRegistry.get_global_platform_registry()
+    
+    PrettyOutput.print(f"Starting Jarvis API server on {host}:{port}", OutputType.SUCCESS)
+    PrettyOutput.print("This server provides an OpenAI-compatible API", OutputType.INFO)
+    
+    if default_platform and default_model:
+        PrettyOutput.print(f"Default platform: {default_platform}, model: {default_model}", OutputType.INFO)
+    
+    PrettyOutput.print("Available platforms:", OutputType.INFO)
+    
+    # Print available platforms and models
+    platforms = registry.get_available_platforms()
+    list_platforms()
+    
+    # Platform and model cache
+    platform_instances = {}
+    
+    # Chat history storage
+    chat_histories = {}
+    
+    def get_platform_instance(platform_name: str, model_name: str):
+        """Get or create a platform instance"""
+        key = f"{platform_name}:{model_name}"
+        if key not in platform_instances:
+            platform = registry.create_platform(platform_name)
+            if not platform:
+                raise HTTPException(status_code=400, detail=f"Platform {platform_name} not found")
+            
+            platform.set_model_name(model_name)
+            platform.set_suppress_output(True)  # Suppress console output in server mode
+            platform_instances[key] = platform
+        
+        return platform_instances[key]
+    
+    def log_conversation(conversation_id, messages, model, response=None):
+        """Log conversation to file in plain text format."""
+        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        log_file = os.path.join(logs_dir, f"conversation_{conversation_id}_{timestamp}.txt")
+
+        with open(log_file, "w", encoding="utf-8") as f:
+            f.write(f"Conversation ID: {conversation_id}\n")
+            f.write(f"Timestamp: {timestamp}\n")
+            f.write(f"Model: {model}\n\n")
+            f.write("Messages:\n")
+            for message in messages:
+                f.write(f"{message['role']}: {message['content']}\n")
+            if response:
+                f.write(f"\nResponse:\n{response}\n")
+        
+        PrettyOutput.print(f"Conversation logged to {log_file}", OutputType.INFO)
+    
+    @app.get("/v1/models")
+    async def list_models():
+        """List available models for the specified platform in OpenAI-compatible format"""
+        model_list = []
+        
+        # Only get models for the currently set platform
+        if default_platform:
+            try:
+                platform = registry.create_platform(default_platform)
+                if platform:
+                    models = platform.get_model_list()
+                    if models:
+                        for model_name, _ in models:
+                            full_name = f"{default_platform}/{model_name}"
+                            model_list.append({
+                                "id": full_name,
+                                "object": "model",
+                                "created": int(time.time()),
+                                "owned_by": default_platform
+                            })
+            except Exception as e:
+                print(f"Error getting models for {default_platform}: {str(e)}")
+        
+        # Return model list
+        return {"object": "list", "data": model_list}
+    
+    @app.post("/v1/chat/completions")
+    @app.options("/v1/chat/completions")  # 添加 OPTIONS 方法支持
+    async def create_chat_completion(request: ChatCompletionRequest):
+        """Create a chat completion in OpenAI-compatible format"""
+        model = request.model
+        messages = request.messages
+        stream = request.stream
+        
+        # Generate a conversation ID if this is a new conversation
+        conversation_id = str(uuid.uuid4())
+        
+        # Extract platform and model name
+        if "/" in model:
+            platform_name, model_name = model.split("/", 1)
+        else:
+            # Use default platform and model if not specified
+            if default_platform and default_model:
+                platform_name, model_name = default_platform, default_model
+            else:
+                platform_name, model_name = "oyi", model  # Default to OYI platform
+        
+        # Get platform instance
+        platform = get_platform_instance(platform_name, model_name)
+        
+        # Convert messages to text format for the platform
+        message_text = ""
+        for msg in messages:
+            role = msg.role
+            content = msg.content
+            
+            if role == "system":
+                message_text += f"System: {content}\n\n"
+            elif role == "user":
+                message_text += f"User: {content}\n\n"
+            elif role == "assistant":
+                message_text += f"Assistant: {content}\n\n"
+        
+        # Store messages in chat history
+        chat_histories[conversation_id] = {
+            "model": model,
+            "messages": [{"role": m.role, "content": m.content} for m in messages]
+        }
+        
+        # Log the conversation
+        log_conversation(conversation_id, 
+                         [{"role": m.role, "content": m.content} for m in messages], 
+                         model)
+        
+        if stream:
+            # Return streaming response
+            return StreamingResponse(
+                stream_chat_response(platform, message_text, model),
+                media_type="text/event-stream"
+            )
+        else:
+            # Get chat response
+            try:
+                response_text = platform.chat_until_success(message_text)
+                
+                # Create response in OpenAI format
+                completion_id = f"chatcmpl-{str(uuid.uuid4())}"
+                
+                # Update chat history with response
+                if conversation_id in chat_histories:
+                    chat_histories[conversation_id]["messages"].append({
+                        "role": "assistant",
+                        "content": response_text
+                    })
+                
+                # Log the conversation with response
+                log_conversation(conversation_id, 
+                                 chat_histories[conversation_id]["messages"], 
+                                 model,
+                                 response_text)
+                
+                return {
+                    "id": completion_id,
+                    "object": "chat.completion",
+                    "created": int(time.time()),
+                    "model": model,
+                    "choices": [
+                        {
+                            "index": 0,
+                            "message": {
+                                "role": "assistant",
+                                "content": response_text
+                            },
+                            "finish_reason": "stop"
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": len(message_text) // 4,  # Rough estimate
+                        "completion_tokens": len(response_text) // 4,  # Rough estimate
+                        "total_tokens": (len(message_text) + len(response_text)) // 4  # Rough estimate
+                    }
+                }
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+    
+    async def stream_chat_response(platform, message, model_name):
+        """Stream chat response in OpenAI-compatible format"""
+        import time
+        import json
+        import uuid
+        from datetime import datetime
+        import os
+        
+        completion_id = f"chatcmpl-{str(uuid.uuid4())}"
+        created_time = int(time.time())
+        conversation_id = str(uuid.uuid4())
+        
+        # Create logs directory if it doesn't exist
+        logs_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs")
+        os.makedirs(logs_dir, exist_ok=True)
+        
+        # 修改第一个yield语句的格式
+        initial_data = {
+            'id': completion_id,
+            'object': 'chat.completion.chunk',
+            'created': created_time,
+            'model': model_name,
+            'choices': [{
+                'index': 0,
+                'delta': {'role': 'assistant'},
+                'finish_reason': None
+            }]
+        }
+        res = json.dumps(initial_data)
+        yield f"data: {res}\n\n"
+        
+        try:
+            # 直接获取聊天响应，而不是尝试捕获stdout
+            response = platform.chat_until_success(message)
+            
+            # 记录完整响应
+            full_response = ""
+            
+            # 如果有响应，将其分块发送
+            if response:
+                # 分成小块以获得更好的流式体验
+                chunk_size = 4  # 每个块的字符数
+                for i in range(0, len(response), chunk_size):
+                    chunk = response[i:i+chunk_size]
+                    full_response += chunk
+                    
+                    # 创建并发送块
+                    chunk_data = {
+                        'id': completion_id,
+                        'object': 'chat.completion.chunk',
+                        'created': created_time,
+                        'model': model_name,
+                        'choices': [{
+                            'index': 0,
+                            'delta': {'content': chunk},
+                            'finish_reason': None
+                        }]
+                    }
+                    
+                    yield f"data: {json.dumps(chunk_data)}\n\n"
+                    
+                    # 小延迟以模拟流式传输
+                    await asyncio.sleep(0.01)
+            else:
+                # 如果没有输出，发送一个空内容块
+                chunk_data = {
+                    'id': completion_id,
+                    'object': 'chat.completion.chunk',
+                    'created': created_time,
+                    'model': model_name,
+                    'choices': [{
+                        'index': 0,
+                        'delta': {'content': "No response from model."},
+                        'finish_reason': None
+                    }]
+                }
+                yield f"data: {json.dumps(chunk_data)}\n\n"
+                full_response = "No response from model."
+            
+            # 修改最终yield语句的格式
+            final_data = {
+                'id': completion_id,
+                'object': 'chat.completion.chunk',
+                'created': created_time,
+                'model': model_name,
+                'choices': [{
+                    'index': 0,
+                    'delta': {},
+                    'finish_reason': 'stop'
+                }]
+            }
+            yield f"data: {json.dumps(final_data)}\n\n"
+            
+            # 发送[DONE]标记
+            yield "data: [DONE]\n\n"
+            
+            # 记录对话到文件
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            log_file = os.path.join(logs_dir, f"stream_conversation_{conversation_id}_{timestamp}.json")
+            
+            log_data = {
+                "conversation_id": conversation_id,
+                "timestamp": timestamp,
+                "model": model_name,
+                "message": message,
+                "response": full_response
+            }
+            
+            with open(log_file, "w", encoding="utf-8") as f:
+                json.dump(log_data, f, ensure_ascii=False, indent=2)
+            
+            PrettyOutput.print(f"Stream conversation logged to {log_file}", OutputType.INFO)
+            
+        except Exception as e:
+            # 发送错误消息
+            error_msg = f"Error: {str(e)}"
+            print(f"Streaming error: {error_msg}")
+
+            res = json.dumps({
+                'id': completion_id,
+                'object': 'chat.completion.chunk',
+                'created': created_time,
+                'model': model_name,
+                'choices': [{
+                    'index': 0,
+                    'delta': {'content': error_msg},
+                    'finish_reason': 'stop'
+                }]
+            })
+            yield f"data: {res}\n\n"
+            yield f"data: {json.dumps({'error': {'message': error_msg, 'type': 'server_error'}})}\n\n"
+            yield "data: [DONE]\n\n"
+            
+            # 记录错误到文件
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            log_file = os.path.join(logs_dir, f"stream_error_{conversation_id}_{timestamp}.json")
+            
+            log_data = {
+                "conversation_id": conversation_id,
+                "timestamp": timestamp,
+                "model": model_name,
+                "message": message,
+                "error": error_msg
+            }
+            
+            with open(log_file, "w", encoding="utf-8") as f:
+                json.dump(log_data, f, ensure_ascii=False, indent=2)
+            
+            PrettyOutput.print(f"Stream error logged to {log_file}", OutputType.ERROR)
+    
+    # Run the server
+    uvicorn.run(app, host=host, port=port)
+
 def main():
     """Main function"""
     import argparse
@@ -121,12 +528,21 @@ def main():
     chat_parser.add_argument('--platform', '-p', help='指定要使用的平台')
     chat_parser.add_argument('--model', '-m', help='指定要使用的模型')
     
+    # service subcommand
+    service_parser = subparsers.add_parser('service', help='启动OpenAI兼容的API服务')
+    service_parser.add_argument('--host', default='127.0.0.1', help='服务主机地址 (默认: 127.0.0.1)')
+    service_parser.add_argument('--port', type=int, default=8000, help='服务端口 (默认: 8000)')
+    service_parser.add_argument('--platform', '-p', help='指定默认平台，当客户端未指定平台时使用')
+    service_parser.add_argument('--model', '-m', help='指定默认模型，当客户端未指定模型时使用')
+    
     args = parser.parse_args()
     
     if args.command == 'info':
         info_command(args)
     elif args.command == 'chat':
         chat_command(args)
+    elif args.command == 'service':
+        service_command(args)
     else:
         parser.print_help()
 

jarvis/jarvis_platform_manager/openai_test.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+"""
+Test script for Jarvis OpenAI-compatible API service.
+"""
+
+import argparse
+import os
+import sys
+from openai import OpenAI
+
+def test_chat(api_base, model, stream=False, interactive=False):
+    """Test chat completion with the API."""
+    client = OpenAI(
+        api_key="dummy-key",  # Not actually used by our service
+        base_url=f"{api_base}/v1"
+    )
+    
+    print(f"Testing chat with model: {model}, stream={stream}")
+    print("=" * 50)
+    
+    try:
+        # First, list available models
+        print("Available models:")
+        models = client.models.list()
+        for m in models.data:
+            print(f"  - {m.id}")
+        print()
+        
+        if interactive:
+            # Interactive chat mode
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."}
+            ]
+            
+            print("Interactive chat mode. Type 'exit' to quit.")
+            print("=" * 50)
+            
+            while True:
+                # Get user input
+                user_input = input("You: ")
+                if user_input.lower() in ['exit', 'quit', 'bye']:
+                    break
+                
+                # Add user message to history
+                messages.append({"role": "user", "content": user_input})
+                
+                # Get response
+                print("Assistant: ", end="", flush=True)
+                
+                if stream:
+                    response = client.chat.completions.create(
+                        model=model,
+                        messages=messages, # type: ignore
+                        stream=True
+                    ) # type: ignore
+                    
+                    # Process the streaming response
+                    assistant_response = ""
+                    for chunk in response:
+                        if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                            content = chunk.choices[0].delta.content
+                            assistant_response += content
+                            print(content, end="", flush=True)
+                    print()
+                else:
+                    response = client.chat.completions.create(
+                        model=model,
+                        messages=messages # type: ignore
+                    )
+                    assistant_response = response.choices[0].message.content
+                    print(assistant_response)
+                
+                # Add assistant response to history
+                messages.append({"role": "assistant", "content": assistant_response}) # type: ignore
+                print()
+            
+            print("=" * 50)
+            print("Chat session ended.")
+            
+        else:
+            # Single request mode
+            print("Sending chat request...")
+            messages = [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello! Tell me a short joke."}
+            ]
+            
+            if stream:
+                print("Response (streaming):")
+                
+                # Use the OpenAI client for streaming
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=messages, # type: ignore
+                    stream=True
+                ) # type: ignore
+                
+                # Process the streaming response
+                full_content = ""
+                for chunk in response:
+                    if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
+                        content = chunk.choices[0].delta.content
+                        full_content += content
+                        print(content, end="", flush=True)
+                
+                print("\n")
+                print(f"Full response: {full_content}")
+            else:
+                print("Response:")
+                response = client.chat.completions.create(
+                    model=model,
+                    messages=messages # type: ignore
+                )
+                print(response.choices[0].message.content)
+        
+        print("=" * 50)
+        print("Test completed successfully!")
+        
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return 1
+    
+    return 0
+
+def main():
+    parser = argparse.ArgumentParser(description="Test Jarvis OpenAI-compatible API")
+    parser.add_argument("--api-base", default="http://localhost:8000", help="API base URL")
+    parser.add_argument("--model", default="gpt-3.5-turbo", help="Model to test (default: gpt-3.5-turbo)")
+    parser.add_argument("--stream", action="store_true", help="Test streaming mode")
+    parser.add_argument("--interactive", "-i", action="store_true", help="Interactive chat mode")
+    
+    args = parser.parse_args()
+    
+    return test_chat(args.api_base, args.model, args.stream, args.interactive)
+
+if __name__ == "__main__":
+    sys.exit(main())

jarvis/jarvis_rag/main.py

@@ -3,8 +3,8 @@ import numpy as np
 import faiss
 from typing import List, Tuple, Optional, Dict
 import pickle
-from jarvis.utils import OutputType, PrettyOutput, get_context_token_count, get_embedding, get_embedding_batch, get_file_md5, get_max_token_count, get_max_paragraph_length, get_min_paragraph_length, get_thread_count, init_gpu_config, load_embedding_model
-from jarvis.utils import init_env
+from jarvis.jarvis_utils import OutputType, PrettyOutput, get_context_token_count, get_embedding, get_embedding_batch, get_file_md5, get_max_token_count, get_max_paragraph_length, get_min_paragraph_length, get_thread_count, init_gpu_config, load_embedding_model
+from jarvis.jarvis_utils import init_env
 from dataclasses import dataclass
 from tqdm import tqdm
 import fitz  # PyMuPDF for PDF files

jarvis/jarvis_smart_shell/main.py

@@ -8,7 +8,7 @@ from yaspin import yaspin # type: ignore
 from yaspin.spinners import Spinners # type: ignore
 
 from jarvis.jarvis_platform.registry import PlatformRegistry
-from jarvis.utils import PrettyOutput, OutputType, get_shell_name, init_env
+from jarvis.jarvis_utils import PrettyOutput, OutputType, get_multiline_input, get_shell_name, init_env
 
 def execute_command(command: str) -> None:
     """Show command and allow user to edit, then execute, Ctrl+C to cancel"""
@@ -70,19 +70,14 @@ Remember: Only return the command itself, without any additional content.
         prefix = f"Current path: {current_path}\n"
         prefix += f"Current shell: {shell}\n"
         
-        # 使用yaspin显示Thinking状态
-        with yaspin(Spinners.dots, text="Thinking", color="yellow") as spinner:
-            # 处理请求
-            result = model.chat_until_success(prefix + request)
-            
-            # 提取命令
-            if result and isinstance(result, str):
-                command = result.strip()
-                spinner.ok("✓")
-                return command
-            
-            spinner.fail("✗")
-            return None
+        result = model.chat_until_success(prefix + request)
+        
+        # 提取命令
+        if result and isinstance(result, str):
+            command = result.strip()
+            return command
+        
+        return None
         
     except Exception as e:
         PrettyOutput.print(f"处理请求失败: {str(e)}", OutputType.WARNING)
@@ -101,15 +96,21 @@ Example:
   %(prog)s "Find documents modified in the last week"
 """)
     
-    # 添加参数
+    # 修改为可选参数，添加从stdin读取的支持
     parser.add_argument(
         "request",
-        help="描述您想要执行的操作, 用自然语言描述"
+        nargs='?',  # 设置为可选参数
+        help="描述您想要执行的操作（用自然语言描述），如果未提供则从标准输入读取"
     )
     
     # 解析参数
     args = parser.parse_args()
     
+    # 添加标准输入处理
+    if not args.request:
+        # 检查是否在交互式终端中运行
+        args.request = get_multiline_input(tip="请输入您要执行的功能：")
+    
     # 处理请求
     command = process_request(args.request)
     

jarvis/jarvis_tools/ask_codebase.py

@@ -1,5 +1,5 @@
 from typing import Dict, Any
-from jarvis.utils import OutputType, PrettyOutput, dont_use_local_model, find_git_root
+from jarvis.jarvis_utils import OutputType, PrettyOutput, dont_use_local_model, find_git_root
 from jarvis.jarvis_codebase.main import CodeBase
 
 class AskCodebaseTool:

jarvis/jarvis_tools/ask_user.py

@@ -1,6 +1,6 @@
 from typing import Dict, Any
 from jarvis.jarvis_tools.base import Tool
-from jarvis.utils import get_multiline_input, PrettyOutput, OutputType
+from jarvis.jarvis_utils import get_multiline_input, PrettyOutput, OutputType
 
 class AskUserTool:
     name="ask_user"