| from __future__ import annotations |
|
|
| import ast |
| import ipaddress |
| import operator |
| import os |
| import re |
| import socket |
| import threading |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any |
| from urllib.parse import urljoin, urlparse |
|
|
| from dotenv import load_dotenv |
|
|
|
|
| load_dotenv() |
|
|
| PERSISTENT_ROOT = Path(os.getenv("PERSISTENT_ROOT", "/data")).expanduser() |
| try: |
| PERSISTENT_ROOT.mkdir(parents=True, exist_ok=True) |
| except OSError: |
| PERSISTENT_ROOT = Path("data") |
| PERSISTENT_ROOT.mkdir(parents=True, exist_ok=True) |
|
|
| os.environ.setdefault("HF_HOME", str(PERSISTENT_ROOT / "huggingface")) |
| os.environ.setdefault("HF_HUB_CACHE", str(PERSISTENT_ROOT / "huggingface" / "hub")) |
|
|
| import gradio as gr |
| import requests |
| from bs4 import BeautifulSoup |
| from ddgs import DDGS |
| from huggingface_hub import hf_hub_download |
| from llama_cpp import Llama |
| from smolagents import ChatMessage, CodeAgent, Model, Tool |
|
|
|
|
| MODEL_REPO = os.getenv("MODEL_REPO", "NANI-Nithin/K2-Horizon-0.9B-GGUF") |
| MODEL_FILE = os.getenv("MODEL_FILE", "K2-Horizon-0.9B-Q4_K_M.gguf") |
| MODEL_DIR = Path(os.getenv("MODEL_DIR", str(PERSISTENT_ROOT / "models"))).expanduser() |
|
|
|
|
| def env_int(name: str, default: int) -> int: |
| try: |
| return int(os.getenv(name, str(default))) |
| except ValueError: |
| return default |
|
|
|
|
| def env_float(name: str, default: float) -> float: |
| try: |
| return float(os.getenv(name, str(default))) |
| except ValueError: |
| return default |
|
|
|
|
| def content_to_text(content: Any) -> str: |
| if isinstance(content, str): |
| return content |
| if isinstance(content, list): |
| parts: list[str] = [] |
| for item in content: |
| if isinstance(item, dict): |
| parts.append(str(item.get("text", item.get("content", "")))) |
| else: |
| parts.append(str(item)) |
| return "\n".join(part for part in parts if part) |
| return str(content or "") |
|
|
|
|
| class LlamaCppModel(Model): |
| def __init__(self, llama: Llama, max_tokens: int, temperature: float) -> None: |
| super().__init__() |
| self.llama = llama |
| self.max_tokens = max_tokens |
| self.temperature = temperature |
|
|
| @staticmethod |
| def _normalize_messages(messages: list[Any]) -> list[dict[str, str]]: |
| normalized: list[dict[str, str]] = [] |
| for message in messages: |
| if isinstance(message, dict): |
| role = message.get("role", "user") |
| content = message.get("content", "") |
| else: |
| role = getattr(message, "role", "user") |
| content = getattr(message, "content", "") |
| if hasattr(role, "value"): |
| role = role.value |
| role = str(role).lower() |
| if role not in {"system", "user", "assistant"}: |
| role = "user" |
| normalized.append({"role": role, "content": content_to_text(content)}) |
| return normalized |
|
|
| def generate( |
| self, |
| messages: list[Any], |
| stop_sequences: list[str] | None = None, |
| response_format: dict[str, Any] | None = None, |
| tools_to_call_from: list[Tool] | None = None, |
| **kwargs: Any, |
| ) -> ChatMessage: |
| del response_format, tools_to_call_from |
| result = self.llama.create_chat_completion( |
| messages=self._normalize_messages(messages), |
| max_tokens=int(kwargs.get("max_tokens", self.max_tokens)), |
| temperature=float(kwargs.get("temperature", self.temperature)), |
| top_p=float(kwargs.get("top_p", 0.9)), |
| repeat_penalty=float(kwargs.get("repeat_penalty", 1.1)), |
| stop=stop_sequences or None, |
| ) |
| content = result["choices"][0]["message"].get("content", "") |
| return ChatMessage(role="assistant", content=content) |
|
|
| def __call__(self, messages: list[Any], **kwargs: Any) -> ChatMessage: |
| return self.generate(messages, **kwargs) |
|
|
| def direct_chat(self, messages: list[dict[str, str]]) -> str: |
| result = self.llama.create_chat_completion( |
| messages=messages, |
| max_tokens=self.max_tokens, |
| temperature=self.temperature, |
| top_p=0.9, |
| repeat_penalty=1.1, |
| ) |
| return str(result["choices"][0]["message"].get("content", "")).strip() |
|
|
|
|
| class DuckDuckGoSearchTool(Tool): |
| name = "web_search" |
| description = "Search the public web with DuckDuckGo. Use it for current facts and external information." |
| inputs = { |
| "query": {"type": "string", "description": "A focused web search query."}, |
| "max_results": { |
| "type": "integer", |
| "description": "Number of results from 1 to 8.", |
| "nullable": True, |
| }, |
| } |
| output_type = "string" |
|
|
| def forward(self, query: str, max_results: int | None = None) -> str: |
| limit = max(1, min(int(max_results or 5), 8)) |
| results = list(DDGS().text(query, max_results=limit)) |
| if not results: |
| return "No search results found." |
| rows = [] |
| for index, item in enumerate(results, 1): |
| title = item.get("title", "Untitled") |
| url = item.get("href", item.get("url", "")) |
| body = item.get("body", "") |
| rows.append(f"{index}. {title}\nURL: {url}\nSnippet: {body}") |
| return "\n\n".join(rows) |
|
|
|
|
| def ensure_public_url(url: str) -> str: |
| parsed = urlparse(url) |
| if parsed.scheme not in {"http", "https"} or not parsed.hostname: |
| raise ValueError("Only public http/https URLs are allowed.") |
| addresses = socket.getaddrinfo(parsed.hostname, parsed.port or 80, proto=socket.IPPROTO_TCP) |
| for address in addresses: |
| ip = ipaddress.ip_address(address[4][0]) |
| if not ip.is_global: |
| raise ValueError("Private, loopback, and local network addresses are blocked.") |
| return url |
|
|
|
|
| class ReadWebpageTool(Tool): |
| name = "read_webpage" |
| description = "Download and extract readable text from a public web page URL." |
| inputs = { |
| "url": {"type": "string", "description": "The full public http or https URL."}, |
| } |
| output_type = "string" |
|
|
| def forward(self, url: str) -> str: |
| current_url = url |
| response = None |
| for _ in range(4): |
| safe_url = ensure_public_url(current_url) |
| response = requests.get( |
| safe_url, |
| timeout=12, |
| allow_redirects=False, |
| headers={"User-Agent": "Mozilla/5.0 (compatible; K2-Horizon-Agent/1.0)"}, |
| ) |
| if response.status_code not in {301, 302, 303, 307, 308}: |
| break |
| location = response.headers.get("location") |
| if not location: |
| break |
| current_url = urljoin(current_url, location) |
| assert response is not None |
| response.raise_for_status() |
| content_type = response.headers.get("content-type", "") |
| if "text/html" not in content_type and "text/plain" not in content_type: |
| return f"Unsupported content type: {content_type}" |
| soup = BeautifulSoup(response.text[:2_000_000], "html.parser") |
| for node in soup(["script", "style", "noscript", "svg"]): |
| node.decompose() |
| text = re.sub(r"\n{3,}", "\n\n", soup.get_text("\n", strip=True)) |
| return text[:12_000] or "No readable text found." |
|
|
|
|
| _BINARY_OPERATORS = { |
| ast.Add: operator.add, |
| ast.Sub: operator.sub, |
| ast.Mult: operator.mul, |
| ast.Div: operator.truediv, |
| ast.FloorDiv: operator.floordiv, |
| ast.Mod: operator.mod, |
| ast.Pow: operator.pow, |
| } |
| _UNARY_OPERATORS = {ast.UAdd: operator.pos, ast.USub: operator.neg} |
|
|
|
|
| def evaluate_expression(node: ast.AST) -> float | int: |
| if isinstance(node, ast.Expression): |
| return evaluate_expression(node.body) |
| if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): |
| return node.value |
| if isinstance(node, ast.BinOp) and type(node.op) in _BINARY_OPERATORS: |
| left = evaluate_expression(node.left) |
| right = evaluate_expression(node.right) |
| if isinstance(node.op, ast.Pow) and abs(right) > 100: |
| raise ValueError("Exponent is too large.") |
| return _BINARY_OPERATORS[type(node.op)](left, right) |
| if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPERATORS: |
| return _UNARY_OPERATORS[type(node.op)](evaluate_expression(node.operand)) |
| raise ValueError("Only numeric arithmetic is supported.") |
|
|
|
|
| class CalculatorTool(Tool): |
| name = "calculator" |
| description = "Safely evaluate a numeric arithmetic expression." |
| inputs = {"expression": {"type": "string", "description": "Arithmetic expression to evaluate."}} |
| output_type = "string" |
|
|
| def forward(self, expression: str) -> str: |
| if len(expression) > 200: |
| raise ValueError("Expression is too long.") |
| value = evaluate_expression(ast.parse(expression, mode="eval")) |
| return str(value) |
|
|
|
|
| class CurrentTimeTool(Tool): |
| name = "current_time" |
| description = "Get the current system date and time, including timezone." |
| inputs = {} |
| output_type = "string" |
|
|
| def forward(self) -> str: |
| return datetime.now().astimezone().isoformat(timespec="seconds") |
|
|
|
|
| class AppRuntime: |
| def __init__(self) -> None: |
| self.lock = threading.Lock() |
| self.model: LlamaCppModel | None = None |
| self.agent: CodeAgent | None = None |
| self.model_path: Path | None = None |
|
|
| def load(self) -> None: |
| if self.model is not None: |
| return |
| with self.lock: |
| if self.model is not None: |
| return |
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
| local_path = hf_hub_download( |
| repo_id=MODEL_REPO, |
| filename=MODEL_FILE, |
| local_dir=str(MODEL_DIR), |
| ) |
| self.model_path = Path(local_path) |
| llama = Llama( |
| model_path=str(self.model_path), |
| n_ctx=env_int("N_CTX", 4096), |
| n_threads=env_int("N_THREADS", max(1, (os.cpu_count() or 4) - 1)), |
| n_threads_batch=env_int("N_THREADS_BATCH", os.cpu_count() or 4), |
| n_batch=env_int("N_BATCH", 256), |
| n_gpu_layers=0, |
| use_mmap=True, |
| verbose=os.getenv("LLAMA_VERBOSE", "0") == "1", |
| ) |
| self.model = LlamaCppModel( |
| llama=llama, |
| max_tokens=env_int("MAX_NEW_TOKENS", 700), |
| temperature=env_float("TEMPERATURE", 0.2), |
| ) |
| self.agent = CodeAgent( |
| tools=[ |
| DuckDuckGoSearchTool(), |
| ReadWebpageTool(), |
| CalculatorTool(), |
| CurrentTimeTool(), |
| ], |
| model=self.model, |
| max_steps=env_int("AGENT_MAX_STEPS", 5), |
| add_base_tools=False, |
| additional_authorized_imports=[], |
| code_block_tags="markdown", |
| ) |
|
|
| def reply(self, message: str, history: list[dict[str, str]], use_tools: bool) -> str: |
| self.load() |
| assert self.model is not None |
| if use_tools: |
| assert self.agent is not None |
| transcript = "\n".join( |
| f"{item.get('role', 'user')}: {content_to_text(item.get('content', ''))}" |
| for item in history[-6:] |
| if item.get("role") in {"user", "assistant"} |
| ) |
| task = message |
| if transcript: |
| task = f"Conversation context:\n{transcript}\n\nCurrent user request:\n{message}" |
| return str(self.agent.run(task, reset=True)).strip() |
| system = { |
| "role": "system", |
| "content": "You are K2 Horizon, a concise and helpful local assistant.", |
| } |
| context = [system] |
| for item in history[-10:]: |
| if item.get("role") in {"user", "assistant"}: |
| context.append({"role": item["role"], "content": content_to_text(item.get("content", ""))}) |
| context.append({"role": "user", "content": message}) |
| return self.model.direct_chat(context) |
|
|
|
|
| runtime = AppRuntime() |
|
|
|
|
| def respond(message: str, history: list[dict[str, str]], use_tools: bool): |
| if not message.strip(): |
| return "", history |
| updated = history + [{"role": "user", "content": message}] |
| try: |
| answer = runtime.reply(message.strip(), history, use_tools) |
| except Exception as exc: |
| answer = f"Error: {type(exc).__name__}: {exc}" |
| updated.append({"role": "assistant", "content": answer}) |
| return "", updated |
|
|
|
|
| CSS = """ |
| .gradio-container { max-width: 860px !important; margin: 0 auto !important; } |
| #app-shell { min-height: 100vh; padding: 32px 12px 20px; } |
| #title { text-align: center; margin-bottom: 2px; } |
| #subtitle { text-align: center; color: var(--body-text-color-subdued); margin-bottom: 18px; } |
| #chat { border: 1px solid var(--border-color-primary); border-radius: 18px; overflow: hidden; } |
| #composer { gap: 10px; align-items: stretch; margin-top: 12px; } |
| #prompt textarea { border-radius: 14px !important; } |
| #send { min-width: 92px; border-radius: 14px !important; } |
| #controls { align-items: center; margin-top: 8px; } |
| #note { color: var(--body-text-color-subdued); font-size: 12px; text-align: right; } |
| footer { display: none !important; } |
| """ |
|
|
| THEME = gr.themes.Base(primary_hue="slate", neutral_hue="slate") |
|
|
|
|
| with gr.Blocks( |
| title="K2 Horizon", |
| ) as demo: |
| with gr.Column(elem_id="app-shell"): |
| gr.Markdown("# K2 Horizon", elem_id="title") |
| gr.Markdown("Private CPU inference with optional web tools", elem_id="subtitle") |
| chatbot = gr.Chatbot( |
| height=570, |
| buttons=["copy"], |
| allow_tags=False, |
| placeholder="Ask anything", |
| elem_id="chat", |
| ) |
| with gr.Row(elem_id="composer"): |
| prompt = gr.Textbox( |
| placeholder="Message K2 Horizon…", |
| show_label=False, |
| scale=9, |
| lines=1, |
| max_lines=5, |
| elem_id="prompt", |
| ) |
| send = gr.Button("Send", variant="primary", scale=1, elem_id="send") |
| with gr.Row(elem_id="controls"): |
| use_tools = gr.Checkbox(value=True, label="Web tools", scale=1) |
| clear = gr.Button("Clear", variant="secondary", size="sm", scale=0) |
| gr.Markdown("Model loads on the first message", elem_id="note") |
|
|
| send.click(respond, [prompt, chatbot, use_tools], [prompt, chatbot]) |
| prompt.submit(respond, [prompt, chatbot, use_tools], [prompt, chatbot]) |
| clear.click(lambda: ("", []), outputs=[prompt, chatbot], queue=False) |
|
|
|
|
| if __name__ == "__main__": |
| demo.queue(default_concurrency_limit=1).launch( |
| server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"), |
| server_port=env_int("GRADIO_SERVER_PORT", 7860), |
| share=os.getenv("GRADIO_SHARE", "0") == "1", |
| show_error=True, |
| ssr_mode=False, |
| footer_links=[], |
| theme=THEME, |
| css=CSS, |
| ) |
|
|
