File size: 2,852 Bytes
5a067eb
 
 
 
7eeb3bc
5a067eb
 
7eeb3bc
 
 
 
5a067eb
 
7eeb3bc
5a067eb
 
 
 
e30c7c9
 
 
 
5a067eb
 
 
e30c7c9
 
 
7eeb3bc
 
6c53eaa
e30c7c9
 
 
7eeb3bc
e30c7c9
6c53eaa
7eeb3bc
 
e30c7c9
 
f4b8dff
 
 
 
 
 
 
 
 
5a067eb
f4b8dff
e30c7c9
f4b8dff
e30c7c9
f4b8dff
e30c7c9
 
6c53eaa
e30c7c9
 
7eeb3bc
f4b8dff
 
 
 
 
e30c7c9
7eeb3bc
2da0450
 
7eeb3bc
5a067eb
 
d8ebbd9
5a067eb
7eeb3bc
5a067eb
 
 
 
7eeb3bc
 
5a067eb
7eeb3bc
 
 
 
 
5a067eb
 
7eeb3bc
 
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
import gradio as gr
import os
import re
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ.get("OPENROUTER_API_KEY"),
)

def clean_response(text: str) -> str:
    lines = []
    for line in text.splitlines():
        line = re.sub(r"^\s*[*#•\-]+\s*", "", line)
        line = re.sub(r"\*\*(.*?)\*\*", r"\1", line)
        # Only strip double underscores around a phrase. A single token like
        # __init__ is a Python dunder, which shares the exact shape of
        # __underline__ and cannot be told apart by position alone.
        line = re.sub(r"(?<!\w)__([^_\n]*?\s[^_\n]*?)__(?!\w)", r"\1", line)
        lines.append(line)
    return "\n".join(lines)

MAX_SENT_TURNS = 20


def chat(user_message: str, history: list):
    history = history or []

    # Cap what is sent, not what is shown: unbounded history eventually exceeds
    # the model's context window, and the failure then looks like a model error.
    recent = history[-(MAX_SENT_TURNS * 2):]
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for msg in recent:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_message})

    history = history + [{"role": "user", "content": user_message}]

    try:
        # openrouter/free automatically picks a working free model
        stream = client.chat.completions.create(
            model="openrouter/free",
            messages=messages,
            stream=True,
            temperature=0.7,
            max_tokens=4096,
        )

        response_text = ""
        history = history + [{"role": "assistant", "content": ""}]
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                response_text += chunk.choices[0].delta.content
                history[-1]["content"] = response_text
                yield history

        history[-1]["content"] = clean_response(response_text)
        yield history

    except Exception as e:
        history.append({
            "role": "assistant",
            "content": f"Error: {str(e)}"
        })
        yield history

with gr.Blocks(title="OpenRouter Chat") as demo:
    gr.Markdown("## OpenRouter Chat")
    
    chatbot = gr.Chatbot(
        label="Conversation",
        height=600
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Ask something",
            placeholder="Type your prompt and press Enter",
            scale=4,
            container=False
        )
        clear = gr.Button("Clear", scale=1)

    msg.submit(chat, [msg, chatbot], chatbot).then(
        lambda: "", None, msg
    )
    clear.click(lambda: [], None, chatbot)

if __name__ == "__main__":
    demo.launch()