Update main.py

main.py

@@ -1,103 +1,107 @@
-import os
-import time
-import random
-import asyncio
+import json
+import sseclient
 import requests
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
-from typing import List, Optional, Union
-
-app = FastAPI()
-
-class ChatCompletionMessage(BaseModel):
-    role: str
-    content: str
-
-class ChatCompletionRequest(BaseModel):
-    model: str
-    messages: List[ChatCompletionMessage]
-    temperature: Optional[float] = 1.0
-    max_tokens: Optional[int] = None
-    stream: Optional[bool] = False
+from flask import Flask, request, Response, stream_with_context
+import random
 
-class ChatCompletionResponse(BaseModel):
-    id: str
-    object: str
-    created: int
-    model: str
-    choices: List[dict]
-    usage: dict
+app = Flask(__name__)
 
 def generate_random_ip():
     return f"{random.randint(1,255)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(0,255)}"
 
-async def fetch_response(messages: List[ChatCompletionMessage], model: str):
-    your_api_url = "https://chatpro.ai-pro.org/api/ask/openAI"
-    headers = {
-        "content-type": "application/json",
-        "X-Forwarded-For": generate_random_ip(),
-        "origin": "https://chatpro.ai-pro.org",
-        "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36"
+def format_openai_response(content, finish_reason=None):
+    return {
+        "id": "chatcmpl-123",
+        "object": "chat.completion.chunk",
+        "created": 1677652288,
+        "model": "gpt-4o",
+        "choices": [{
+            "delta": {"content": content} if content else {"finish_reason": finish_reason},
+            "index": 0,
+            "finish_reason": finish_reason
+        }]
     }
+
+@app.route('/hf/v1/chat/completions', methods=['POST'])
+def chat_completions():
+    data = request.json
+    messages = data.get('messages', [])
+    stream = data.get('stream', False)
+    
+    if not messages:
+        return {"error": "No messages provided"}, 400
     
     # 将消息列表转换为单个字符串，保留对话历史
-    conversation = "\n".join([f"{msg.role}: {msg.content}" for msg in messages])
+    conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
     
     # 添加指导语
     conversation += "\n请关注并回复user最近的消息并避免总结对话历史的回答"
-    
-    data = {
+
+    model = data.get('model', 'gpt-4o')
+
+    original_api_url = 'https://chatpro.ai-pro.org/api/ask/openAI'
+    headers = {
+        'content-type': 'application/json',
+        'X-Forwarded-For': generate_random_ip(),
+        'origin': 'https://chatpro.ai-pro.org',
+        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36'
+    }
+    payload = {
         "text": conversation,
         "endpoint": "openAI",
         "model": model
     }
 
-    response = requests.post(your_api_url, headers=headers, json=data)
-
-    if response.status_code != 200:
-        raise HTTPException(status_code=response.status_code, detail="Error from upstream API")
-
-    return response.json()
-
-async def stream_response(content: str):
-    # Send the entire content as a single chunk
-    yield f"data: {{'id': 'chatcmpl-{os.urandom(12).hex()}', 'object': 'chat.completion.chunk', 'created': 1677652288, 'model': 'gpt-3.5-turbo-0613', 'choices': [{'index': 0, 'delta': {{'content': '{content}'}}, 'finish_reason': None}]}}\n\n"
-    yield f"data: {{'id': 'chatcmpl-{os.urandom(12).hex()}', 'object': 'chat.completion.chunk', 'created': 1677652288, 'model': 'gpt-3.5-turbo-0613', 'choices': [{'index': 0, 'delta': {{}}, 'finish_reason': 'stop'}]}}\n\n"
-    yield 'data: [DONE]\n\n'
+    def generate():
+        last_content = ""
+        response = requests.post(original_api_url, headers=headers, json=payload, stream=True)
+        client = sseclient.SSEClient(response)
+        
+        for event in client.events():
+            if event.data.startswith('{"text":'):
+                data = json.loads(event.data)
+                new_content = data['text'][len(last_content):]
+                last_content = data['text']
+                
+                if new_content:
+                    yield f"data: {json.dumps(format_openai_response(new_content))}\n\n"
+            
+            elif '"final":true' in event.data:
+                yield f"data: {json.dumps(format_openai_response('', 'stop'))}\n\n"
+                yield "data: [DONE]\n\n"
+                break
 
-@app.post("/hf/v1/chat/completions")
-async def chat_completions(request: Request):
-    body = await request.json()
-    chat_request = ChatCompletionRequest(**body)
-
-    # 传递整个消息历史到API
-    api_response = await fetch_response(chat_request.messages, chat_request.model)
-
-    content = api_response.get("response", "")
-
-    if chat_request.stream:
-        return StreamingResponse(stream_response(content), media_type="text/event-stream")
+    if stream:
+        return Response(stream_with_context(generate()), content_type='text/event-stream')
     else:
-        openai_response = ChatCompletionResponse(
-            id="chatcmpl-" + os.urandom(12).hex(),
-            object="chat.completion",
-            created=int(time.time()),
-            model=chat_request.model,
-            choices=[
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": content
-                    },
-                    "finish_reason": "stop"
-                }
-            ],
-            usage={
-                "prompt_tokens": sum(len(msg.content) for msg in chat_request.messages),
-                "completion_tokens": len(content),
-                "total_tokens": sum(len(msg.content) for msg in chat_request.messages) + len(content)
+        full_response = ""
+        for chunk in generate():
+            if chunk.startswith("data: ") and not chunk.strip() == "data: [DONE]":
+                response_data = json.loads(chunk[6:])
+                if 'choices' in response_data and response_data['choices']:
+                    delta = response_data['choices'][0].get('delta', {})
+                    if 'content' in delta:
+                        full_response += delta['content']
+
+        return {
+            "id": "chatcmpl-123",
+            "object": "chat.completion",
+            "created": 1677652288,
+            "model": model,
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": full_response
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
             }
-        )
-        return openai_response
+        }
+
+if __name__ == '__main__':
+    app.run(debug=True, port=5000)