modularized interface and added tools

ui/inference.py

@@ -1,42 +0,0 @@
-# ui/inference.py
-from huggingface_hub import InferenceClient
-
-import gradio as gr
-
-MODEL_ID = "Qwen/Qwen3.6-27B"
-
-
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model=MODEL_ID)
-
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response

ui/inference/__init__.py

@@ -0,0 +1,4 @@
+# ui/inference/__init__.py
+from .respond import respond
+
+__all__ = ["respond"]

ui/inference/completion.py

@@ -0,0 +1,88 @@
+# ui/inference/completion.py
+from __future__ import annotations
+
+from typing import Any
+
+from huggingface_hub import InferenceClient
+
+from .messages import parse_text_tool_calls, parse_tool_calls
+
+
+def complete_turn(
+    client: InferenceClient,
+    api_messages: list[dict[str, Any]],
+    *,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+    tools: list[dict[str, Any]] | None,
+    tool_choice: str = "auto",
+) -> tuple[str, str, list[dict[str, Any]] | None]:
+    """Run one model turn, preferring streaming and falling back to a single request."""
+    content = ""
+    reasoning = ""
+    tool_calls_map: dict[int, dict[str, Any]] = {}
+
+    stream = client.chat_completion(
+        api_messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        tools=tools,
+        tool_choice=tool_choice if tools else None,
+        stream=True,
+    )
+    for chunk in stream:
+        if not chunk.choices:
+            continue
+        delta = chunk.choices[0].delta
+        if delta.content:
+            content += delta.content
+        if delta.reasoning:
+            reasoning += delta.reasoning
+        if delta.tool_calls:
+            for tool_call in delta.tool_calls:
+                idx = tool_call.index
+                if idx not in tool_calls_map:
+                    tool_calls_map[idx] = {
+                        "id": tool_call.id,
+                        "type": tool_call.type,
+                        "function": {
+                            "name": tool_call.function.name or "",
+                            "arguments": tool_call.function.arguments or "",
+                        },
+                    }
+                    continue
+                if tool_call.function.name:
+                    tool_calls_map[idx]["function"]["name"] = tool_call.function.name
+                if tool_call.function.arguments:
+                    tool_calls_map[idx]["function"]["arguments"] += (
+                        tool_call.function.arguments
+                    )
+
+    text_tool_calls = parse_text_tool_calls(content) or parse_text_tool_calls(reasoning)
+    tool_calls = list(tool_calls_map.values()) or text_tool_calls
+    if content or reasoning or tool_calls:
+        if text_tool_calls:
+            content = ""
+            reasoning = ""
+        return content, reasoning, tool_calls
+
+    response = client.chat_completion(
+        api_messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        tools=tools,
+        tool_choice=tool_choice if tools else None,
+    )
+    assistant_msg = response.choices[0].message
+    content = assistant_msg.content or ""
+    reasoning = assistant_msg.reasoning or ""
+    text_tool_calls = parse_text_tool_calls(content) or parse_text_tool_calls(reasoning)
+    tool_calls = parse_tool_calls(assistant_msg) or text_tool_calls
+    return (
+        "" if text_tool_calls else content,
+        "" if text_tool_calls else reasoning,
+        tool_calls,
+    )

ui/inference/config.py

@@ -0,0 +1,36 @@
+# ui/inference/config.py
+from __future__ import annotations
+
+from typing import Any
+
+MODEL_ID = "Qwen/Qwen3.6-27B"
+MAX_TOOL_ROUNDS = 5
+
+TOOLS: list[dict[str, Any]] = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_economic_indicator",
+            "description": (
+                "Fetch World Bank economic indicator time series for a country. "
+                "Use ISO-3 country codes (e.g. IND for India, HKG, USA, CHN) and "
+                "World Bank indicator codes (e.g. NY.GDP.MKTP.CD for GDP in "
+                "current US dollars)."
+            ),
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "indicator": {
+                        "type": "string",
+                        "description": "World Bank indicator code, e.g. NY.GDP.MKTP.CD",
+                    },
+                    "country": {
+                        "type": "string",
+                        "description": "ISO-3 country code, e.g. IND, HKG, USA",
+                    },
+                },
+                "required": ["indicator", "country"],
+            },
+        },
+    },
+]

ui/inference/messages.py

@@ -0,0 +1,110 @@
+# ui/inference/messages.py
+from __future__ import annotations
+
+import json
+import re
+import uuid
+from typing import Any
+
+from .config import TOOLS
+
+
+TOOL_NAMES = {tool["function"]["name"] for tool in TOOLS}
+
+
+def history_to_api_messages(history: list[dict[str, Any]]) -> list[dict[str, str]]:
+    """Convert Gradio chat history to API messages, skipping UI-only tool steps."""
+    messages: list[dict[str, str]] = []
+    for msg in history:
+        if msg.get("metadata"):
+            continue
+        content = msg.get("content")
+        role = msg.get("role")
+        if role in ("user", "assistant") and isinstance(content, str) and content:
+            messages.append({"role": role, "content": content})
+    return messages
+
+
+def assistant_message_dict(
+    content: str,
+    tool_calls: list[dict[str, Any]] | None,
+) -> dict[str, Any]:
+    message: dict[str, Any] = {"role": "assistant", "content": content}
+    if tool_calls:
+        message["tool_calls"] = tool_calls
+    return message
+
+
+def parse_tool_calls(message: Any) -> list[dict[str, Any]] | None:
+    if not message.tool_calls:
+        return None
+    return [
+        {
+            "id": tool_call.id,
+            "type": tool_call.type,
+            "function": {
+                "name": tool_call.function.name,
+                "arguments": tool_call.function.arguments,
+            },
+        }
+        for tool_call in message.tool_calls
+    ]
+
+
+def parse_text_tool_calls(text: str) -> list[dict[str, Any]] | None:
+    """Recover tool calls when a model emits them as text instead of tool_calls."""
+    if not text:
+        return None
+
+    for tool_name in TOOL_NAMES:
+        arguments = _parse_json_call(text, tool_name) or _parse_python_call(
+            text, tool_name
+        )
+        if arguments is None:
+            continue
+        return [
+            {
+                "id": f"call_{uuid.uuid4().hex}",
+                "type": "function",
+                "function": {
+                    "name": tool_name,
+                    "arguments": json.dumps(arguments),
+                },
+            }
+        ]
+
+    return None
+
+
+def _parse_json_call(text: str, tool_name: str) -> dict[str, Any] | None:
+    match = re.search(rf"(?:default_api:)?{re.escape(tool_name)}\s*\{{", text)
+    if not match:
+        return None
+
+    start = text.find("{", match.start())
+    if start == -1:
+        return None
+
+    decoder = json.JSONDecoder()
+    try:
+        arguments, _ = decoder.raw_decode(text[start:])
+    except json.JSONDecodeError:
+        return _parse_key_values(text[start:])
+
+    return arguments if isinstance(arguments, dict) else None
+
+
+def _parse_python_call(text: str, tool_name: str) -> dict[str, Any] | None:
+    match = re.search(rf"{re.escape(tool_name)}\s*\((?P<args>[^)]*)\)", text)
+    if not match:
+        return None
+    return _parse_key_values(match.group("args"))
+
+
+def _parse_key_values(text: str) -> dict[str, str] | None:
+    pairs = re.findall(r'(\w+)\s*=\s*["\']([^"\']+)["\']', text)
+    if not pairs:
+        pairs = re.findall(r'["\'](\w+)["\']\s*:\s*["\']([^"\']+)["\']', text)
+    if not pairs:
+        return None
+    return {key: value for key, value in pairs}

ui/inference/respond.py

@@ -0,0 +1,112 @@
+# ui/inference/respond.py
+from __future__ import annotations
+
+from typing import Any
+
+import gradio as gr
+from gradio import ChatMessage
+from huggingface_hub import InferenceClient
+
+from .completion import complete_turn
+from .config import MAX_TOOL_ROUNDS, MODEL_ID, TOOLS
+from .messages import history_to_api_messages
+from .streaming import yield_response
+from .tools import execute_tool_calls
+
+
+def respond(
+    message,
+    history: list[dict[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    hf_token: gr.OAuthToken,
+):
+    """
+    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+    """
+    client = InferenceClient(token=hf_token.token, model=MODEL_ID)
+
+    api_messages: list[dict[str, Any]] = [
+        {"role": "system", "content": system_message},
+        *history_to_api_messages(history),
+        {"role": "user", "content": message},
+    ]
+    ui_messages: list[ChatMessage] = []
+
+    try:
+        for _ in range(MAX_TOOL_ROUNDS):
+            content, reasoning, tool_calls = complete_turn(
+                client,
+                api_messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                tools=TOOLS,
+                tool_choice="auto",
+            )
+
+            if tool_calls:
+                yield from execute_tool_calls(
+                    api_messages, ui_messages, tool_calls, content
+                )
+                continue
+
+            answer = content or reasoning
+            if answer:
+                yield from yield_response(ui_messages, answer)
+                return
+
+            # Retry once with a required tool call when the model returned nothing.
+            content, reasoning, tool_calls = complete_turn(
+                client,
+                api_messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                tools=TOOLS,
+                tool_choice="required",
+            )
+            if tool_calls:
+                yield from execute_tool_calls(
+                    api_messages, ui_messages, tool_calls, content
+                )
+                continue
+
+            answer = content or reasoning
+            if answer:
+                yield from yield_response(ui_messages, answer)
+                return
+
+            # Some providers return an empty first turn when tools are enabled.
+            content, reasoning, _ = complete_turn(
+                client,
+                api_messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                tools=None,
+            )
+            answer = content or reasoning
+            if answer:
+                yield from yield_response(ui_messages, answer)
+                return
+
+            yield from yield_response(
+                ui_messages,
+                "I could not generate a response. Please try again.",
+            )
+            return
+
+    except Exception as exc:
+        yield from yield_response(
+            ui_messages,
+            f"Sorry, something went wrong while generating a response: {exc}",
+        )
+        return
+
+    yield from yield_response(
+        ui_messages,
+        "I reached the maximum number of tool calls for this request.",
+    )

ui/inference/streaming.py

@@ -0,0 +1,32 @@
+# ui/inference/streaming.py
+from __future__ import annotations
+
+from gradio import ChatMessage
+
+
+def yield_streaming_string(text: str):
+    if not text:
+        yield ""
+        return
+    for index in range(len(text)):
+        yield text[: index + 1]
+
+
+def yield_streaming_messages(ui_messages: list[ChatMessage], text: str):
+    messages = list(ui_messages)
+    if not text:
+        messages.append(ChatMessage(role="assistant", content=""))
+        yield messages
+        return
+
+    messages.append(ChatMessage(role="assistant", content=""))
+    for index in range(len(text)):
+        messages[-1] = ChatMessage(role="assistant", content=text[: index + 1])
+        yield messages
+
+
+def yield_response(ui_messages: list[ChatMessage], text: str):
+    if ui_messages:
+        yield from yield_streaming_messages(ui_messages, text)
+    else:
+        yield from yield_streaming_string(text)

ui/inference/tools.py

@@ -0,0 +1,89 @@
+# ui/inference/tools.py
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+from gradio import ChatMessage
+
+from apis.world_bank import get_indicator_series
+
+from .messages import assistant_message_dict
+
+
+def truncate(text: str, limit: int = 4000) -> str:
+    if len(text) <= limit:
+        return text
+    return text[:limit] + "\n… (truncated)"
+
+
+def run_tool(name: str, arguments: str) -> str:
+    if name != "get_economic_indicator":
+        return json.dumps({"error": f"Unknown tool: {name}"})
+
+    try:
+        args = json.loads(arguments or "{}")
+        observations = get_indicator_series(
+            indicator=args["indicator"],
+            country=args["country"],
+        )
+        return json.dumps(observations, default=str)
+    except Exception as exc:
+        return json.dumps({"error": str(exc)})
+
+
+def execute_tool_calls(
+    api_messages: list[dict[str, Any]],
+    ui_messages: list[ChatMessage],
+    tool_calls: list[dict[str, Any]],
+    content: str,
+):
+    api_messages.append(assistant_message_dict(content, tool_calls))
+
+    for tool_call in tool_calls:
+        tool_name = tool_call["function"]["name"]
+        tool_args = tool_call["function"]["arguments"]
+        started = time.monotonic()
+
+        ui_messages.append(
+            ChatMessage(
+                role="assistant",
+                content=(
+                    f"Calling `{tool_name}` with arguments:\n"
+                    f"```json\n{tool_args}\n```"
+                ),
+                metadata={
+                    "title": f"🛠️ Used tool {tool_name}",
+                    "status": "pending",
+                },
+            )
+        )
+        yield ui_messages
+
+        result = run_tool(tool_name, tool_args)
+        duration = time.monotonic() - started
+
+        ui_messages[-1] = ChatMessage(
+            role="assistant",
+            content=(
+                f"Calling `{tool_name}` with arguments:\n"
+                f"```json\n{tool_args}\n```\n\n"
+                f"Result:\n```json\n{truncate(result)}\n```"
+            ),
+            metadata={
+                "title": f"🛠️ Used tool {tool_name}",
+                "status": "done",
+                "duration": duration,
+            },
+        )
+        yield ui_messages
+
+        api_messages.append(
+            {
+                "role": "tool",
+                "tool_call_id": tool_call["id"],
+                "name": tool_name,
+                "content": result,
+            }
+        )