bridge_server.py

# /home/wangxinxin/my_graphrag/bridge_server.py
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import uvicorn
import httpx
import time
import json
import re
import asyncio

app = FastAPI()

# Configure your actual model endpoint
LLM_URL = "http://10.119.70.11:8088/generate"

def sanitize_json_string(s):
    """
    Core repair function:
    Iterates through the string. If unescaped newlines or tabs are found 
    inside a JSON string value, they are forcibly converted to \\n or \\t.
    """
    new_s = []
    in_string = False
    is_escaped = False
    
    for c in s:
        if is_escaped:
            new_s.append(c)
            is_escaped = False
            continue
            
        if c == '\\':
            is_escaped = True
            new_s.append(c)
            continue
            
        if c == '"':
            in_string = not in_string
            new_s.append(c)
            continue
            
        if in_string:
            if c == '\n':
                new_s.append('\\n')
            elif c == '\t':
                new_s.append('\\t')
            elif c == '\r':
                pass # Discard carriage return
            else:
                new_s.append(c)
        else:
            new_s.append(c)
            
    return "".join(new_s)

def extract_json(text):
    """
    Enhanced JSON Extractor v2:
    1. Extracts JSON blocks.
    2. Fixes illegal control characters (resolves "Invalid control character" errors).
    """
    text = text.strip()
    
    # 1. Try direct parsing
    try:
        json.loads(text)
        return text
    except:
        pass

    # 2. Try extracting from markdown blocks: ```json ... ``` or ``` ... ```
    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
    if match:
        text = match.group(1)
    else:
        # 3. Brute-force extract the outermost {}
        start = text.find('{')
        end = text.rfind('}')
        if start != -1 and end != -1 and end > start:
            text = text[start:end+1]

    # 4. Try parsing again; if it fails, perform "cleaning/repair"
    try:
        json.loads(text)
        return text
    except:
        # Key step: call the repair function to handle newlines/tabs
        print("[Bridge] Invalid JSON detected, attempting to fix control characters...")
        fixed_text = sanitize_json_string(text)
        try:
            json.loads(fixed_text)
            print("[Bridge] Repair successful!")
            return fixed_text
        except:
            # If it still fails, return the original snippet and hope for the best
            print("[Bridge] Repair failed, returning original snippet.")
            return text

async def fake_stream_generator(content, model_name):
    """Mimics OpenAI format streaming output for the full text content"""
    chunk_role = {
        "id": "chatcmpl-stream",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model_name,
        "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}]
    }
    yield f"data: {json.dumps(chunk_role)}\n\n"
    
    chunk_content = {
        "id": "chatcmpl-stream",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model_name,
        "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}]
    }
    yield f"data: {json.dumps(chunk_content)}\n\n"
    
    chunk_stop = {
        "id": "chatcmpl-stream",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model_name,
        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
    }
    yield f"data: {json.dumps(chunk_stop)}\n\n"
    yield "data: [DONE]\n\n"

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        data = await request.json()
        messages = data.get("messages", [])
        stream = data.get("stream", False)
        
        prompt = ""
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        prompt += "<|im_start|>assistant\n"

        # --- Determine if JSON cleaning is needed ---
        prompt_lower = prompt.lower()
        # Drift Search prompts usually contain "output as json" or similar instructions
        is_json_request = "json" in prompt_lower
        
        print(f"\n[Request] Stream={stream} | JSON_Mode={is_json_request}")

        # --- CALL LLM ---
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": 2048,
                "temperature": data.get("temperature", 0.3),
                "stop": ["<|im_end|>"],
                "details": True
            }
        }

        async with httpx.AsyncClient(timeout=180.0) as client:
            resp = await client.post(LLM_URL, json=payload)
            resp.raise_for_status()
            result = resp.json()

        raw_text = result[0].get("generated_text", "")
        print(f"[DEBUG] LLM Output : {raw_text}")
        # --- PROCESSING ---
        # Apply cleaning logic if it's a JSON request or looks like an extraction task
        if is_json_request or "extract" in prompt_lower:
            final_content = extract_json(raw_text)
        else:
            final_content = raw_text

        # --- RESPONSE ---
        if stream:
            return StreamingResponse(
                fake_stream_generator(final_content, "sensenova"),
                media_type="text/event-stream"
            )
        else:
            return {
                "id": "chatcmpl-custom",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "sensenova",
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": final_content
                    },
                    "finish_reason": "stop"
                }],
                "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
            }

    except Exception as e:
        print(f"[ERROR] {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8900)