File size: 3,078 Bytes
852bc31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f1e93d
 
852bc31
 
4f1e93d
 
 
 
 
 
 
 
 
 
852bc31
 
 
4f1e93d
 
 
852bc31
4f1e93d
852bc31
 
 
 
 
 
 
 
 
 
 
4f1e93d
852bc31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4f1e93d
852bc31
 
 
4f1e93d
852bc31
 
4f1e93d
852bc31
 
 
4f1e93d
852bc31
 
4f1e93d
852bc31
 
 
 
 
 
 
 
 
4f1e93d
852bc31
 
 
 
 
 
 
 
 
 
 
 
4f1e93d
852bc31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import asyncio
import time
import uuid
import json

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
from gradio_client import Client

app = FastAPI()

# HuggingFace Space
client = Client("CohereLabs/command-a-vision")


# ✅ FIXED: call gradio with positional args
def call_gradio(message, max_tokens=100):

    try:
        # format input like Gradio expects
        payload = {
            "text": message,
            "files": []
        }

        # IMPORTANT: positional inputs (NOT keyword args)
        result = client.predict(
            payload,       # input 1
            max_tokens,    # input 2
            api_name="/chat"
        )

        # result comes as dict sometimes
        if isinstance(result, dict):
            return json.dumps(result)

        return str(result)

    except Exception as e:
        print("Gradio API error:", e)
        return "Error: upstream model failed."


def format_openai_response(content):
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "command-a-vision",
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": content
                },
                "finish_reason": "stop"
            }
        ]
    }


@app.post("/v1/chat/completions")
async def chat(request: Request):

    body = await request.json()

    messages = body.get("messages", [])
    stream = body.get("stream", False)

    max_tokens = body.get("max_tokens", 100)

    user_message = messages[-1]["content"]

    # ✅ normal response
    if not stream:

        result = call_gradio(user_message, max_tokens)

        return JSONResponse(format_openai_response(result))

    # ✅ streaming response
    async def generate():

        result = call_gradio(user_message, max_tokens)

        words = result.split(" ")

        for word in words:

            chunk = {
                "id": f"chatcmpl-{uuid.uuid4().hex}",
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": "command-a-vision",
                "choices": [
                    {
                        "delta": {"content": word + " "},
                        "index": 0,
                        "finish_reason": None
                    }
                ]
            }

            yield f"data: {json.dumps(chunk)}\n\n"
            await asyncio.sleep(0.02)

        # end
        end_chunk = {
            "id": f"chatcmpl-{uuid.uuid4().hex}",
            "object": "chat.completion.chunk",
            "choices": [
                {
                    "delta": {},
                    "index": 0,
                    "finish_reason": "stop"
                }
            ]
        }

        yield f"data: {json.dumps(end_chunk)}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")