|
|
from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect |
|
|
from fastapi.responses import JSONResponse, HTMLResponse |
|
|
from fastapi.staticfiles import StaticFiles |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from pydantic import BaseModel, Field |
|
|
import uvicorn |
|
|
import json |
|
|
import os |
|
|
import asyncio |
|
|
from datetime import datetime |
|
|
from typing import List, Dict, Optional, Any |
|
|
import logging |
|
|
import uuid |
|
|
|
|
|
|
|
|
try: |
|
|
from recursive_thinking_ai import EnhancedRecursiveThinkingChat |
|
|
except ModuleNotFoundError: |
|
|
|
|
|
import sys |
|
|
sys.path.append('.') |
|
|
from recursive_thinking_ai import EnhancedRecursiveThinkingChat |
|
|
|
|
|
|
|
|
logging.basicConfig( |
|
|
level=logging.INFO, |
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' |
|
|
) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
app = FastAPI( |
|
|
title="Chain-of-Recursive-Thoughts: TEST", |
|
|
description="https://github.com/PhialsBasement/Chain-of-Recursive-Thoughts", |
|
|
version="1.0.0" |
|
|
) |
|
|
|
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
allow_credentials=True, |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
|
|
|
chat_instances = {} |
|
|
|
|
|
|
|
|
API_KEY = os.getenv("OPENROUTE_API") |
|
|
if not API_KEY: |
|
|
logger.warning("The OPENROUTE_API environment variable is not set. Some features may not work.") |
|
|
|
|
|
|
|
|
class ChatConfig(BaseModel): |
|
|
model: str = "mistralai/mistral-small-3.1-24b-instruct:free" |
|
|
temperature: Optional[float] = Field(default=0.7, ge=0.0, le=1.0) |
|
|
|
|
|
class MessageRequest(BaseModel): |
|
|
session_id: str |
|
|
message: str |
|
|
thinking_rounds: Optional[int] = Field(default=None, ge=1, le=10) |
|
|
alternatives_per_round: Optional[int] = Field(default=3, ge=1, le=5) |
|
|
temperature: Optional[float] = Field(default=None, ge=0.0, le=1.0) |
|
|
|
|
|
class SaveRequest(BaseModel): |
|
|
session_id: str |
|
|
filename: Optional[str] = None |
|
|
full_log: bool = False |
|
|
|
|
|
class SessionInfo(BaseModel): |
|
|
session_id: str |
|
|
message_count: int |
|
|
created_at: str |
|
|
model: str |
|
|
|
|
|
class SessionResponse(BaseModel): |
|
|
sessions: List[SessionInfo] |
|
|
|
|
|
class InitResponse(BaseModel): |
|
|
session_id: str |
|
|
status: str |
|
|
|
|
|
@app.get("/", response_class=HTMLResponse) |
|
|
async def root(): |
|
|
""" |
|
|
Root endpoint with a simple HTML interface. |
|
|
Modified to display side-by-side "Original" vs. "Chain-of-Thought" outputs, |
|
|
and to show real-time "thinking in progress" messages. |
|
|
""" |
|
|
html_content = f""" |
|
|
<!DOCTYPE html> |
|
|
<html> |
|
|
<head> |
|
|
<title>Chain-of-Recursive-Thoughts: TEST</title> |
|
|
<style> |
|
|
body {{ |
|
|
font-family: Arial, sans-serif; |
|
|
max-width: 1000px; |
|
|
margin: 0 auto; |
|
|
padding: 20px; |
|
|
line-height: 1.6; |
|
|
}} |
|
|
h1 {{ |
|
|
color: #333; |
|
|
border-bottom: 1px solid #eee; |
|
|
padding-bottom: 10px; |
|
|
}} |
|
|
.container {{ |
|
|
background-color: #f9f9f9; |
|
|
border-radius: 5px; |
|
|
padding: 20px; |
|
|
margin-top: 20px; |
|
|
}} |
|
|
label {{ |
|
|
display: block; |
|
|
margin-bottom: 5px; |
|
|
font-weight: bold; |
|
|
}} |
|
|
input, textarea, select {{ |
|
|
width: 100%; |
|
|
padding: 8px; |
|
|
margin-bottom: 10px; |
|
|
border: 1px solid #ddd; |
|
|
border-radius: 4px; |
|
|
}} |
|
|
button {{ |
|
|
background-color: #4CAF50; |
|
|
color: white; |
|
|
padding: 10px 15px; |
|
|
border: none; |
|
|
border-radius: 4px; |
|
|
cursor: pointer; |
|
|
}} |
|
|
button:hover {{ |
|
|
background-color: #45a049; |
|
|
}} |
|
|
#original-response, #chain-response {{ |
|
|
white-space: pre-wrap; |
|
|
background-color: #f5f5f5; |
|
|
padding: 15px; |
|
|
border-radius: 4px; |
|
|
min-height: 100px; |
|
|
}} |
|
|
.responses-container {{ |
|
|
display: flex; |
|
|
gap: 20px; |
|
|
margin-top: 20px; |
|
|
}} |
|
|
.response-column {{ |
|
|
flex: 1; |
|
|
display: flex; |
|
|
flex-direction: column; |
|
|
}} |
|
|
.column-title {{ |
|
|
margin-top: 0; |
|
|
}} |
|
|
.log {{ |
|
|
margin-top: 20px; |
|
|
font-size: 0.9em; |
|
|
color: #666; |
|
|
}} |
|
|
.thinking-progress {{ |
|
|
margin-top: 10px; |
|
|
font-style: italic; |
|
|
color: #555; |
|
|
}} |
|
|
</style> |
|
|
</head> |
|
|
<body> |
|
|
<h1>Chain-of-Recursive-Thoughts: TEST</h1> |
|
|
<div class="container"> |
|
|
<div id="init-form"> |
|
|
<h2>1. Initialize Chat</h2> |
|
|
|
|
|
<label for="model">Model:</label> |
|
|
<input type="text" id="model" value="mistralai/mistral-small-3.1-24b-instruct:free"> |
|
|
|
|
|
<label for="temperature">Temperature:</label> |
|
|
<input type="number" id="temperature" min="0" max="1" step="0.1" value="0.7"> |
|
|
|
|
|
<button onclick="initializeChat()">Initialize</button> |
|
|
</div> |
|
|
|
|
|
<div id="chat-form" style="display: none;"> |
|
|
<h2>2. Send Message</h2> |
|
|
<p>Session ID: <span id="session-id"></span></p> |
|
|
|
|
|
<label for="message">Message:</label> |
|
|
<textarea id="message" rows="4" placeholder="Enter your message"></textarea> |
|
|
|
|
|
<label for="thinking-rounds">Thinking Rounds (optional):</label> |
|
|
<input type="number" id="thinking-rounds" min="1" max="10" placeholder="Auto"> |
|
|
|
|
|
<label for="alternatives">Number of Alternatives (optional):</label> |
|
|
<input type="number" id="alternatives" min="1" max="5" value="3"> |
|
|
|
|
|
<button onclick="sendMessage()">Send</button> |
|
|
<button onclick="resetChat()" style="background-color: #f44336;">Reset</button> |
|
|
</div> |
|
|
|
|
|
<div class="responses-container" id="response-container" style="display: none;"> |
|
|
<!-- Original Response Column --> |
|
|
<div class="response-column"> |
|
|
<h2 class="column-title">Original (No Chain-of-Thought)</h2> |
|
|
<div id="original-response">The original LLM response will appear here...</div> |
|
|
</div> |
|
|
|
|
|
<!-- Chain-of-Thought Response Column --> |
|
|
<div class="response-column"> |
|
|
<h2 class="column-title">Chain-of-Thought</h2> |
|
|
<div id="chain-thinking-progress" class="thinking-progress"></div> |
|
|
<div id="chain-response">The chain-of-thought response will appear here...</div> |
|
|
<div class="log"> |
|
|
<h3>Thinking Process Log:</h3> |
|
|
<div id="thinking-log"></div> |
|
|
</div> |
|
|
</div> |
|
|
</div> |
|
|
</div> |
|
|
|
|
|
<div style="margin-top: 30px;"> |
|
|
<p>Repo: https://github.com/PhialsBasement/Chain-of-Recursive-Thoughts</p> |
|
|
<p>Community: https://discord.gg/openfreeai</p> |
|
|
</div> |
|
|
|
|
|
<script> |
|
|
let currentSessionId = null; |
|
|
let ws = null; // WebSocket for Chain-of-Thought |
|
|
|
|
|
async function initializeChat() {{ |
|
|
const model = document.getElementById('model').value; |
|
|
const temperature = parseFloat(document.getElementById('temperature').value); |
|
|
|
|
|
try {{ |
|
|
const response = await fetch('/api/initialize', {{ |
|
|
method: 'POST', |
|
|
headers: {{ |
|
|
'Content-Type': 'application/json', |
|
|
}}, |
|
|
body: JSON.stringify({{ |
|
|
model: model, |
|
|
temperature: temperature |
|
|
}}), |
|
|
}}); |
|
|
|
|
|
const data = await response.json(); |
|
|
|
|
|
if (response.ok) {{ |
|
|
currentSessionId = data.session_id; |
|
|
document.getElementById('session-id').textContent = currentSessionId; |
|
|
document.getElementById('init-form').style.display = 'none'; |
|
|
document.getElementById('chat-form').style.display = 'block'; |
|
|
document.getElementById('response-container').style.display = 'flex'; |
|
|
|
|
|
// Connect to WebSocket for chain-of-thought streaming |
|
|
connectWebSocket(); |
|
|
}} else {{ |
|
|
alert('Initialization failed: ' + (data.detail || 'Unknown error')); |
|
|
}} |
|
|
}} catch (error) {{ |
|
|
alert('An error occurred: ' + error.message); |
|
|
}} |
|
|
}} |
|
|
|
|
|
function connectWebSocket() {{ |
|
|
if (!currentSessionId) {{ |
|
|
return; |
|
|
}} |
|
|
const loc = window.location; |
|
|
let wsUrl = ''; |
|
|
if (loc.protocol === 'https:') {{ |
|
|
wsUrl = 'wss:'; |
|
|
}} else {{ |
|
|
wsUrl = 'ws:'; |
|
|
}} |
|
|
wsUrl += '//' + loc.host + '/ws/' + currentSessionId; |
|
|
ws = new WebSocket(wsUrl); |
|
|
|
|
|
ws.onopen = function() {{ |
|
|
console.log('WebSocket connection opened for chain-of-thought.'); |
|
|
}}; |
|
|
|
|
|
ws.onmessage = function(event) {{ |
|
|
const data = JSON.parse(event.data); |
|
|
|
|
|
if (data.type === 'status') {{ |
|
|
// Show "thinking in progress" messages |
|
|
document.getElementById('chain-thinking-progress').textContent = data.message; |
|
|
}} else if (data.type === 'chunk') {{ |
|
|
// In this example code, "chunk" is not truly partial, but let's append anyway |
|
|
document.getElementById('chain-thinking-progress').textContent += '\\n' + data.content; |
|
|
}} else if (data.type === 'final') {{ |
|
|
// Final answer |
|
|
document.getElementById('chain-thinking-progress').textContent = ''; |
|
|
document.getElementById('chain-response').textContent = data.response; |
|
|
|
|
|
// Display thinking history |
|
|
const thinkingLog = document.getElementById('thinking-log'); |
|
|
thinkingLog.innerHTML = ''; |
|
|
data.thinking_history.forEach(item => {{ |
|
|
const selected = item.selected ? ' ✓ Selected' : ''; |
|
|
let logEntry = document.createElement('p'); |
|
|
logEntry.innerHTML = '<strong>Round ' + item.round + selected + ':</strong> '; |
|
|
if (item.explanation && item.selected) {{ |
|
|
logEntry.innerHTML += '<br><em>Reason for selection: ' + item.explanation + '</em>'; |
|
|
}} |
|
|
thinkingLog.appendChild(logEntry); |
|
|
}}); |
|
|
}} else if (data.type === 'error') {{ |
|
|
document.getElementById('chain-thinking-progress').textContent = ''; |
|
|
document.getElementById('chain-response').textContent = 'Error: ' + data.error; |
|
|
}} else if (data.error) {{ |
|
|
// Session not found or other errors |
|
|
document.getElementById('chain-thinking-progress').textContent = ''; |
|
|
document.getElementById('chain-response').textContent = 'Error: ' + data.error; |
|
|
}} |
|
|
}}; |
|
|
|
|
|
ws.onerror = function(error) {{ |
|
|
console.error('WebSocket error:', error); |
|
|
document.getElementById('chain-thinking-progress').textContent = 'WebSocket error. Check console.'; |
|
|
}}; |
|
|
|
|
|
ws.onclose = function() {{ |
|
|
console.log('WebSocket closed.'); |
|
|
}}; |
|
|
}} |
|
|
|
|
|
async function sendMessage() {{ |
|
|
if (!currentSessionId) {{ |
|
|
alert('Please initialize a chat session first.'); |
|
|
return; |
|
|
}} |
|
|
|
|
|
const message = document.getElementById('message').value; |
|
|
const thinkingRounds = document.getElementById('thinking-rounds').value; |
|
|
const alternatives = document.getElementById('alternatives').value; |
|
|
|
|
|
if (!message) {{ |
|
|
alert('Please enter a message.'); |
|
|
return; |
|
|
}} |
|
|
|
|
|
// Clear out previous responses/log |
|
|
document.getElementById('original-response').textContent = 'Loading original response...'; |
|
|
document.getElementById('chain-response').textContent = ''; |
|
|
document.getElementById('thinking-log').textContent = ''; |
|
|
document.getElementById('chain-thinking-progress').textContent = 'Thinking...'; |
|
|
|
|
|
// 1) Get Original Response via standard fetch |
|
|
try {{ |
|
|
const originalRes = await fetch('/api/send_message_original', {{ |
|
|
method: 'POST', |
|
|
headers: {{ |
|
|
'Content-Type': 'application/json', |
|
|
}}, |
|
|
body: JSON.stringify({{ |
|
|
session_id: currentSessionId, |
|
|
message: message |
|
|
}}), |
|
|
}}); |
|
|
|
|
|
const originalData = await originalRes.json(); |
|
|
if (originalRes.ok) {{ |
|
|
document.getElementById('original-response').textContent = originalData.response; |
|
|
}} else {{ |
|
|
document.getElementById('original-response').textContent = 'Error: ' + (originalData.detail || 'Unknown error'); |
|
|
}} |
|
|
}} catch (error) {{ |
|
|
document.getElementById('original-response').textContent = 'Error: ' + error.message; |
|
|
}} |
|
|
|
|
|
// 2) Send message to WebSocket for chain-of-thought |
|
|
if (ws && ws.readyState === WebSocket.OPEN) {{ |
|
|
const payload = {{ |
|
|
type: 'message', |
|
|
content: message, |
|
|
thinking_rounds: thinkingRounds ? parseInt(thinkingRounds) : null, |
|
|
alternatives_per_round: alternatives ? parseInt(alternatives) : 3 |
|
|
}}; |
|
|
ws.send(JSON.stringify(payload)); |
|
|
}} else {{ |
|
|
document.getElementById('chain-thinking-progress').textContent = 'WebSocket not connected. Unable to get chain-of-thought response.'; |
|
|
}} |
|
|
}} |
|
|
|
|
|
function resetChat() {{ |
|
|
currentSessionId = null; |
|
|
if (ws) {{ |
|
|
ws.close(); |
|
|
ws = null; |
|
|
}} |
|
|
document.getElementById('init-form').style.display = 'block'; |
|
|
document.getElementById('chat-form').style.display = 'none'; |
|
|
document.getElementById('response-container').style.display = 'none'; |
|
|
document.getElementById('message').value = ''; |
|
|
document.getElementById('thinking-rounds').value = ''; |
|
|
document.getElementById('alternatives').value = '3'; |
|
|
}} |
|
|
</script> |
|
|
</body> |
|
|
</html> |
|
|
""" |
|
|
return HTMLResponse(content=html_content) |
|
|
|
|
|
@app.get("/health") |
|
|
async def health_check(): |
|
|
"""Health check endpoint""" |
|
|
return {"status": "healthy", "timestamp": datetime.now().isoformat()} |
|
|
|
|
|
@app.post("/api/initialize", response_model=InitResponse) |
|
|
async def initialize_chat(config: ChatConfig): |
|
|
"""Initialize a new chat session using the environment API key""" |
|
|
try: |
|
|
|
|
|
session_id = f"session_{datetime.now().strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}" |
|
|
|
|
|
if not API_KEY: |
|
|
raise HTTPException(status_code=400, detail="The OPENROUTE_API environment variable is not set.") |
|
|
|
|
|
chat = EnhancedRecursiveThinkingChat( |
|
|
api_key=API_KEY, |
|
|
model=config.model, |
|
|
temperature=config.temperature |
|
|
) |
|
|
chat_instances[session_id] = { |
|
|
"chat": chat, |
|
|
"created_at": datetime.now().isoformat(), |
|
|
"model": config.model |
|
|
} |
|
|
return {"session_id": session_id, "status": "initialized"} |
|
|
except Exception as e: |
|
|
logger.error(f"Error initializing chat: {str(e)}") |
|
|
raise HTTPException(status_code=500, detail=f"Failed to initialize chat: {str(e)}") |
|
|
|
|
|
@app.post("/api/send_message_original") |
|
|
async def send_message_original(request: MessageRequest): |
|
|
""" |
|
|
Return a direct LLM response without applying the chain-of-thought logic. |
|
|
""" |
|
|
if request.session_id not in chat_instances: |
|
|
raise HTTPException(status_code=404, detail="Session not found") |
|
|
|
|
|
chat = chat_instances[request.session_id]["chat"] |
|
|
try: |
|
|
|
|
|
messages = [{"role": "user", "content": request.message}] |
|
|
response_data = chat._call_api(messages, temperature=chat.temperature, stream=False) |
|
|
|
|
|
|
|
|
|
|
|
if isinstance(response_data, dict) and "choices" in response_data: |
|
|
first_choice = response_data["choices"][0] |
|
|
if "message" in first_choice and "content" in first_choice["message"]: |
|
|
original_text = first_choice["message"]["content"] |
|
|
elif "text" in first_choice: |
|
|
original_text = first_choice["text"] |
|
|
else: |
|
|
original_text = str(first_choice) |
|
|
else: |
|
|
|
|
|
original_text = str(response_data) |
|
|
|
|
|
return {"response": original_text.strip()} |
|
|
except Exception as e: |
|
|
logger.error(f"Error getting original response: {str(e)}") |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
|
|
|
@app.post("/api/send_message") |
|
|
async def send_message(request: MessageRequest): |
|
|
""" |
|
|
Send a message and get a response with the chain-of-thought process (HTTP-based, not streaming). |
|
|
Primarily left here for completeness, but the user-facing code calls the WebSocket for streaming. |
|
|
""" |
|
|
try: |
|
|
if request.session_id not in chat_instances: |
|
|
raise HTTPException(status_code=404, detail="Session not found") |
|
|
|
|
|
chat = chat_instances[request.session_id]["chat"] |
|
|
|
|
|
original_thinking_fn = chat._determine_thinking_rounds |
|
|
original_alternatives_fn = chat._generate_alternatives |
|
|
original_temperature = getattr(chat, "temperature", 0.7) |
|
|
|
|
|
if request.thinking_rounds is not None: |
|
|
chat._determine_thinking_rounds = lambda _: request.thinking_rounds |
|
|
|
|
|
if request.alternatives_per_round is not None: |
|
|
def modified_generate_alternatives(base_response, prompt, num_alternatives=3): |
|
|
return original_alternatives_fn(base_response, prompt, request.alternatives_per_round) |
|
|
chat._generate_alternatives = modified_generate_alternatives |
|
|
|
|
|
if request.temperature is not None: |
|
|
setattr(chat, "temperature", request.temperature) |
|
|
|
|
|
logger.info(f"Processing message for session {request.session_id}") |
|
|
start_time = datetime.now() |
|
|
result = chat.think_and_respond(request.message, verbose=True) |
|
|
processing_time = (datetime.now() - start_time).total_seconds() |
|
|
logger.info(f"Message processed in {processing_time:.2f} seconds") |
|
|
|
|
|
|
|
|
chat._determine_thinking_rounds = original_thinking_fn |
|
|
chat._generate_alternatives = original_alternatives_fn |
|
|
if request.temperature is not None: |
|
|
setattr(chat, "temperature", original_temperature) |
|
|
|
|
|
return { |
|
|
"session_id": request.session_id, |
|
|
"response": result["response"], |
|
|
"thinking_rounds": result["thinking_rounds"], |
|
|
"thinking_history": result["thinking_history"], |
|
|
"processing_time": processing_time |
|
|
} |
|
|
except Exception as e: |
|
|
logger.error(f"Error processing message: {str(e)}") |
|
|
raise HTTPException(status_code=500, detail=f"Failed to process message: {str(e)}") |
|
|
|
|
|
@app.post("/api/save") |
|
|
async def save_conversation(request: SaveRequest): |
|
|
"""Save the conversation or the full thinking log.""" |
|
|
try: |
|
|
if request.session_id not in chat_instances: |
|
|
raise HTTPException(status_code=404, detail="Session not found") |
|
|
|
|
|
chat = chat_instances[request.session_id]["chat"] |
|
|
|
|
|
filename = request.filename |
|
|
if filename is None: |
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
|
|
log_type = "full_log" if request.full_log else "conversation" |
|
|
filename = f"recthink_{log_type}_{timestamp}.json" |
|
|
|
|
|
os.makedirs("logs", exist_ok=True) |
|
|
file_path = os.path.join("logs", filename) |
|
|
|
|
|
if request.full_log: |
|
|
chat.save_full_log(file_path) |
|
|
else: |
|
|
chat.save_conversation(file_path) |
|
|
|
|
|
return {"status": "saved", "filename": filename, "path": file_path} |
|
|
except Exception as e: |
|
|
logger.error(f"Error saving conversation: {str(e)}") |
|
|
raise HTTPException(status_code=500, detail=f"Failed to save conversation: {str(e)}") |
|
|
|
|
|
@app.get("/api/sessions", response_model=SessionResponse) |
|
|
async def list_sessions(): |
|
|
"""List all active chat sessions.""" |
|
|
sessions = [] |
|
|
for session_id, session_data in chat_instances.items(): |
|
|
chat = session_data["chat"] |
|
|
message_count = len(chat.conversation_history) // 2 |
|
|
sessions.append(SessionInfo( |
|
|
session_id=session_id, |
|
|
message_count=message_count, |
|
|
created_at=session_data["created_at"], |
|
|
model=session_data["model"] |
|
|
)) |
|
|
return {"sessions": sessions} |
|
|
|
|
|
@app.get("/api/sessions/{session_id}") |
|
|
async def get_session(session_id: str): |
|
|
"""Get details for a specific chat session.""" |
|
|
if session_id not in chat_instances: |
|
|
raise HTTPException(status_code=404, detail="Session not found") |
|
|
|
|
|
session_data = chat_instances[session_id] |
|
|
chat = session_data["chat"] |
|
|
|
|
|
conversation = [] |
|
|
for i in range(0, len(chat.conversation_history), 2): |
|
|
if i+1 < len(chat.conversation_history): |
|
|
conversation.append({ |
|
|
"user": chat.conversation_history[i], |
|
|
"assistant": chat.conversation_history[i+1] |
|
|
}) |
|
|
|
|
|
return { |
|
|
"session_id": session_id, |
|
|
"created_at": session_data["created_at"], |
|
|
"model": session_data["model"], |
|
|
"message_count": len(conversation), |
|
|
"conversation": conversation |
|
|
} |
|
|
|
|
|
@app.delete("/api/sessions/{session_id}") |
|
|
async def delete_session(session_id: str): |
|
|
"""Delete a chat session.""" |
|
|
if session_id not in chat_instances: |
|
|
raise HTTPException(status_code=404, detail="Session not found") |
|
|
del chat_instances[session_id] |
|
|
return {"status": "deleted", "session_id": session_id} |
|
|
|
|
|
class ConnectionManager: |
|
|
def __init__(self): |
|
|
self.active_connections: Dict[str, WebSocket] = {} |
|
|
|
|
|
async def connect(self, session_id: str, websocket: WebSocket): |
|
|
await websocket.accept() |
|
|
self.active_connections[session_id] = websocket |
|
|
|
|
|
def disconnect(self, session_id: str): |
|
|
if session_id in self.active_connections: |
|
|
del self.active_connections[session_id] |
|
|
|
|
|
async def send_json(self, session_id: str, data: dict): |
|
|
if session_id in self.active_connections: |
|
|
await self.active_connections[session_id].send_json(data) |
|
|
|
|
|
manager = ConnectionManager() |
|
|
|
|
|
@app.websocket("/ws/{session_id}") |
|
|
async def websocket_endpoint(websocket: WebSocket, session_id: str): |
|
|
try: |
|
|
await manager.connect(session_id, websocket) |
|
|
|
|
|
if session_id not in chat_instances: |
|
|
await websocket.send_json({"error": "Session not found"}) |
|
|
await websocket.close() |
|
|
return |
|
|
|
|
|
chat = chat_instances[session_id]["chat"] |
|
|
original_call_api = chat._call_api |
|
|
|
|
|
async def stream_callback(chunk): |
|
|
await manager.send_json(session_id, {"type": "chunk", "content": chunk}) |
|
|
|
|
|
def ws_call_api(messages, temperature=0.7, stream=True): |
|
|
result = original_call_api(messages, temperature, stream) |
|
|
if stream: |
|
|
asyncio.create_task(stream_callback(result)) |
|
|
return result |
|
|
|
|
|
chat._call_api = ws_call_api |
|
|
|
|
|
while True: |
|
|
data = await websocket.receive_text() |
|
|
message_data = json.loads(data) |
|
|
|
|
|
if message_data["type"] == "message": |
|
|
start_time = datetime.now() |
|
|
|
|
|
try: |
|
|
thinking_rounds = message_data.get("thinking_rounds", None) |
|
|
alternatives_per_round = message_data.get("alternatives_per_round", None) |
|
|
temperature = message_data.get("temperature", None) |
|
|
|
|
|
original_thinking_fn = chat._determine_thinking_rounds |
|
|
original_alternatives_fn = chat._generate_alternatives |
|
|
original_temperature = getattr(chat, "temperature", 0.7) |
|
|
|
|
|
if thinking_rounds is not None: |
|
|
chat._determine_thinking_rounds = lambda _: thinking_rounds |
|
|
|
|
|
if alternatives_per_round is not None: |
|
|
def modified_generate_alternatives(base_response, prompt, num_alternatives=3): |
|
|
return original_alternatives_fn(base_response, prompt, alternatives_per_round) |
|
|
chat._generate_alternatives = modified_generate_alternatives |
|
|
|
|
|
if temperature is not None: |
|
|
setattr(chat, "temperature", temperature) |
|
|
|
|
|
await manager.send_json(session_id, { |
|
|
"type": "status", |
|
|
"status": "processing", |
|
|
"message": "Starting recursive thinking process..." |
|
|
}) |
|
|
|
|
|
result = chat.think_and_respond(message_data["content"], verbose=True) |
|
|
processing_time = (datetime.now() - start_time).total_seconds() |
|
|
|
|
|
chat._determine_thinking_rounds = original_thinking_fn |
|
|
chat._generate_alternatives = original_alternatives_fn |
|
|
if temperature is not None: |
|
|
setattr(chat, "temperature", original_temperature) |
|
|
|
|
|
await manager.send_json(session_id, { |
|
|
"type": "final", |
|
|
"response": result["response"], |
|
|
"thinking_rounds": result["thinking_rounds"], |
|
|
"thinking_history": result["thinking_history"], |
|
|
"processing_time": processing_time |
|
|
}) |
|
|
|
|
|
except Exception as e: |
|
|
error_msg = str(e) |
|
|
logger.error(f"Error in WebSocket message processing: {error_msg}") |
|
|
await manager.send_json(session_id, { |
|
|
"type": "error", |
|
|
"error": error_msg |
|
|
}) |
|
|
except WebSocketDisconnect: |
|
|
logger.info(f"WebSocket disconnected: {session_id}") |
|
|
manager.disconnect(session_id) |
|
|
except Exception as e: |
|
|
error_msg = str(e) |
|
|
logger.error(f"WebSocket error: {error_msg}") |
|
|
try: |
|
|
await websocket.send_json({"type": "error", "error": error_msg}) |
|
|
except: |
|
|
pass |
|
|
finally: |
|
|
if 'chat' in locals() and 'original_call_api' in locals(): |
|
|
chat._call_api = original_call_api |
|
|
|
|
|
manager.disconnect(session_id) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
port = 7860 |
|
|
print(f"Starting server on port {port}") |
|
|
uvicorn.run("app:app", host="0.0.0.0", port=port) |
|
|
|
