wot-floweval-core/llmInterface.py at main · SystemWorkflows/wot-floweval-core · GitHub

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
125
126
127
128
129
130
131
132
133
134
135
136
137
from google import genai
from google.genai import types
from openai import OpenAI
from typing import Any, Dict
import time
import json
class llmFactory:
    llmChat = None

    @staticmethod
    def produceChat(name:str, parameters:Dict[str, Any]):
        match name:
            case "models/gemini-2.5-flash":             return geminiChat(parameters)
            case "models/gemini-2.5-pro":               return geminiChat(parameters)
            case "models/gemini-2.5-flash-lite":        return geminiChat(parameters)
            case "openai/gpt-oss-120b":                 return openRouterChat(parameters)
            case "openai/o4-mini-high":                 return openRouterChat(parameters)
            case "openai/gpt-5-mini":                   return openRouterChat(parameters)
            case "openai/gpt-5":                        return openRouterChat(parameters)
            case "openai/gpt-oss-20b:free":             return openRouterChat(parameters)
            case "deepseek/deepseek-r1:free":           return openRouterChat(parameters)
            case "deepseek/deepseek-chat-v3.1:free":    return openRouterChat(parameters)
            case "openrouter/sonoma-sky-alpha":         return openRouterChat(parameters)
            case "z-ai/glm-4.5-air:free":               return openRouterChat(parameters)
            case "qwen/qwen3-coder:free":               return openRouterChat(parameters)
            case _:                                     return None


class llmChat:
    def __init__(self, parameters):
        pass

    def chat(self, prompt):
        pass

class geminiChat(llmChat):
    def __init__(self, parameters):
        client = genai.Client(api_key="AIzaSyCvx6ffmJxe0XkcApWKDIMi5jrX3W7kQ6U") # No billing connected
        self.chat = client.chats.create(
            model=parameters["model"],
            config=types.GenerateContentConfig(
                temperature=parameters["temperature"],
                #candidate_count=parameters["candidate_count"],
                seed=parameters["seed"],
                max_output_tokens=parameters["max_output_tokens"],
                response_mime_type=parameters["response_mime_type"]
            )
        )

    def send_message(self, prompt):
        time_start = time.time()
        response = self.chat.send_message(prompt)
        time_end = time.time()
        delta_t = time_end - time_start
        resp = None

        try:
            resp = response.candidates[0].content.parts[0].text
        except:
            raise Exception("failed to generate flow")

        if resp is None:
            raise Exception("failed to generate flow")

        return {
            "response": resp,
            "metadata":{
                "time": delta_t,
                "input_tokens": response.usage_metadata.prompt_token_count,
                "output_tokens": response.usage_metadata.candidates_token_count
            }
        }


class openRouterChat(llmChat):
    def __init__(self, parameters):
        self.parameters = parameters
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key="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"#no billing connected
        )
        self.messages = []


    def send_message(self, prompt):
        self.messages.append({
        "role": "user",
        "content": [
            {
            "type": "text",
            "text": prompt
            }
        ]
        })
        time_start = time.time()
        completion = self.client.chat.completions.create(
            model=self.parameters.get("model"),
            messages=self.messages,
            temperature=self.parameters["temperature"],
            seed=self.parameters["seed"],
            reasoning_effort="high", # low, medium, high
            extra_body={
                "provider": {
                    "require_parameters": True,
                    "only":[
                        "ncompass"
                    ],
                    "allow_fallbacks": False
                }
            }
        )
        time_end = time.time()
        delta_t = time_end - time_start
        resp = None
        try:
            resp = completion.choices[0].message.content
            if resp == None:
                raise Exception("failed to generate flow")
        except:
            raise Exception("failed to generate flow")

        self.messages.append({
            "role": "assistant",
            "content": [
                {
                "type": "text",
                "text": resp
                }
            ]
        })

        resp = resp.strip()
        if resp[3:7] == "json":
            resp = resp[8:-3]
        resp = resp.strip("`")

        return {"response": resp, "metadata":{"time": delta_t, "input_tokens": completion.usage.prompt_tokens, "output_tokens": completion.usage.completion_tokens}}