asky-cli 0.1.6__py3-none-any.whl

asky/llm.py

@@ -0,0 +1,378 @@
+"""LLM integration and conversation loop."""
+
+import json
+import os
+import re
+import time
+from typing import Any, Dict, List, Optional
+
+import requests
+from rich.console import Console
+from rich.markdown import Markdown
+
+
+from asky.config import (
+    MAX_TURNS,
+    MODELS,
+    TOOLS,
+    SYSTEM_PROMPT,
+    FORCE_SEARCH_PROMPT,
+    SYSTEM_PROMPT_SUFFIX,
+    DEEP_RESEARCH_PROMPT_TEMPLATE,
+    DEEP_DIVE_PROMPT_TEMPLATE,
+    QUERY_SUMMARY_MAX_CHARS,
+    ANSWER_SUMMARY_MAX_CHARS,
+    SUMMARIZATION_MODEL,
+    LLM_USER_AGENT,
+    SUMMARIZE_QUERY_PROMPT_TEMPLATE,
+    SUMMARIZE_ANSWER_PROMPT_TEMPLATE,
+    REQUEST_TIMEOUT,
+    DEFAULT_CONTEXT_SIZE,
+    CONTINUE_QUERY_THRESHOLD,
+)
+from asky.html import strip_think_tags
+from asky.tools import dispatch_tool_call, reset_read_urls
+
+
+def is_markdown(text: str) -> bool:
+    """Check if the text likely contains markdown formatting."""
+    # Basic detection: common markdown patterns
+    patterns = [
+        r"^#+\s",  # Headers
+        r"\*\*.*\*\*",  # Bold
+        r"__.*__",  # Bold
+        r"\*.*\* ",  # Italic
+        r"_.*_",  # Italic
+        r"\[.*\]\(.*\)",  # Links
+        r"```",  # Code blocks
+        r"^\s*[-*+]\s",  # Lists
+        r"^\s*\d+\.\s",  # Numbered lists
+    ]
+    return any(re.search(p, text, re.M) for p in patterns)
+
+
+def parse_textual_tool_call(text: str) -> Optional[Dict[str, Any]]:
+    """Parse tool calls from textual format (fallback for some models)."""
+    if not text:
+        return None
+    m = re.search(r"to=functions\.([a-zA-Z0-9_]+)", text)
+    if not m:
+        return None
+    name = m.group(1)
+    j = re.search(r"(\{.*\})", text, re.S)
+    if not j:
+        return None
+    try:
+        json.loads(j.group(1))
+        return {"name": name, "arguments": j.group(1)}
+    except Exception:
+        return None
+
+
+class UsageTracker:
+    """Track token usage per model alias."""
+
+    def __init__(self):
+        self.usage: Dict[str, int] = {}
+
+    def add_usage(self, model_alias: str, tokens: int):
+        self.usage[model_alias] = self.usage.get(model_alias, 0) + tokens
+
+    def get_total_usage(self, model_alias: str) -> int:
+        return self.usage.get(model_alias, 0)
+
+
+def count_tokens(messages: List[Dict[str, Any]]) -> int:
+    """Naive token counting: chars / 4."""
+    total_chars = 0
+    for m in messages:
+        content = m.get("content")
+        if content:
+            total_chars += len(content)
+        # Also count tool calls and results
+        tc = m.get("tool_calls")
+        if tc:
+            total_chars += len(json.dumps(tc))
+    return total_chars // 4
+
+
+def get_llm_msg(
+    model_id: str,
+    messages: List[Dict[str, Any]],
+    tools: Optional[List[Dict]] = TOOLS,
+    verbose: bool = False,
+    model_alias: Optional[str] = None,
+    usage_tracker: Optional[UsageTracker] = None,
+) -> Dict[str, Any]:
+    """Send messages to the LLM and get a response."""
+    # Find the model config based on model_id
+    model_config = next((m for m in MODELS.values() if m["id"] == model_id), None)
+
+    url = ""
+    headers = {
+        "Content-Type": "application/json",
+        "User-Agent": LLM_USER_AGENT,
+    }
+
+    if model_config and "base_url" in model_config:
+        url = model_config["base_url"]
+
+    api_key = None
+    if model_config and "api_key" in model_config:
+        api_key = model_config["api_key"]
+    elif model_config and "api_key_env" in model_config:
+        api_key_env_var = model_config["api_key_env"]
+        api_key = os.environ.get(api_key_env_var)
+        if not api_key:
+            print(f"Warning: {api_key_env_var} not found in environment variables.")
+
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    payload = {
+        "model": model_id,
+        "messages": messages,
+    }
+    if tools:
+        payload["tools"] = tools
+        payload["tool_choice"] = "auto"
+
+    max_retries = 10
+    backoff = 2
+    max_backoff = 60
+
+    if verbose:
+        print(f"\n[DEBUG] Sending to LLM ({model_id})...")
+        print(f"Tools enabled: {bool(tools)}")
+        print("Last message sent:")
+        if messages:
+            last_msg = messages[-1]
+            content = last_msg.get("content", "")
+            if content and len(content) > 500:
+                print(f"  Role: {last_msg['role']}")
+                print(f"  Content (truncated): {content[:500]}...")
+            else:
+                print(f"  {json.dumps(last_msg, indent=2)}")
+        print("-" * 20)
+
+    tokens_sent = count_tokens(messages)
+    if model_alias:
+        print(f"[{model_alias}] Sent: {tokens_sent} tokens")
+    else:
+        print(f"[{model_id}] Sent: {tokens_sent} tokens")
+
+    for attempt in range(max_retries):
+        try:
+            resp = requests.post(
+                url, json=payload, headers=headers, timeout=REQUEST_TIMEOUT
+            )
+            resp.raise_for_status()
+            resp_json = resp.json()
+
+            # Extract usage if available, otherwise use naive count
+            usage = resp_json.get("usage", {})
+            prompt_tokens = usage.get("prompt_tokens", tokens_sent)
+            completion_tokens = usage.get("completion_tokens", 0)
+            if "completion_tokens" not in usage:
+                completion_tokens = (
+                    len(json.dumps(resp_json["choices"][0]["message"])) // 4
+                )
+
+            total_call_tokens = prompt_tokens + completion_tokens
+
+            if usage_tracker and model_alias:
+                usage_tracker.add_usage(model_alias, total_call_tokens)
+
+            return resp_json["choices"][0]["message"]
+        except requests.exceptions.HTTPError as e:
+            if e.response is not None and e.response.status_code == 429:
+                if attempt < max_retries - 1:
+                    retry_after = e.response.headers.get("Retry-After")
+                    if retry_after:
+                        try:
+                            # Handle potential floating point strings (e.g. "5.0")
+                            wait_time = int(float(retry_after))
+                        except ValueError:
+                            wait_time = backoff
+                            backoff = min(backoff * 2, max_backoff)
+                    else:
+                        wait_time = backoff
+                        backoff = min(backoff * 2, max_backoff)
+
+                    print(
+                        f"Rate limit exceeded (429). Retrying in {wait_time} seconds..."
+                    )
+                    time.sleep(wait_time)
+                    continue
+            raise e
+        except requests.exceptions.RequestException as e:
+            if attempt < max_retries - 1:
+                print(f"Request error: {e}. Retrying in {backoff} seconds...")
+                time.sleep(backoff)
+                backoff = min(backoff * 2, max_backoff)
+                continue
+            raise e
+    raise requests.exceptions.RequestException("Max retries exceeded")
+
+
+def extract_calls(msg: Dict[str, Any], turn: int) -> List[Dict[str, Any]]:
+    """Extract tool calls from an LLM message."""
+    tc = msg.get("tool_calls")
+    if tc:
+        return tc
+    parsed = parse_textual_tool_call(msg.get("content", ""))
+    if parsed:
+        return [{"id": f"textual_call_{turn}", "function": parsed}]
+    return []
+
+
+def construct_system_prompt(
+    deep_research_n: int, deep_dive: bool, force_search: bool
+) -> str:
+    """Build the system prompt based on mode flags."""
+    system_content = SYSTEM_PROMPT
+    if force_search:
+        system_content += FORCE_SEARCH_PROMPT
+    system_content += SYSTEM_PROMPT_SUFFIX
+
+    if deep_research_n > 0:
+        system_content += DEEP_RESEARCH_PROMPT_TEMPLATE.format(n=deep_research_n)
+    if deep_dive:
+        system_content += DEEP_DIVE_PROMPT_TEMPLATE
+    return system_content
+
+
+def run_conversation_loop(
+    model_config: Dict[str, Any],
+    messages: List[Dict[str, Any]],
+    summarize: bool,
+    verbose: bool = False,
+    usage_tracker: Optional[UsageTracker] = None,
+) -> str:
+    """Run the multi-turn conversation loop with tool execution."""
+    turn = 0
+    start_time = time.perf_counter()
+    final_answer = ""
+    original_system_prompt = (
+        messages[0]["content"] if messages and messages[0]["role"] == "system" else ""
+    )
+
+    # Reset read URLs for new conversation
+    reset_read_urls()
+
+    try:
+        while turn < MAX_TURNS:
+            turn += 1
+
+            # Token & Turn Tracking
+            total_tokens = count_tokens(messages)
+            context_size = model_config.get("context_size", DEFAULT_CONTEXT_SIZE)
+            turns_left = MAX_TURNS - turn + 1
+
+            status_msg = (
+                f"\n\n[SYSTEM UPDATE]:\n"
+                f"- Context Used: {total_tokens / context_size * 100:.2f}%"
+                f"- Turns Remaining: {turns_left} (out of {MAX_TURNS})\n"
+                f"Please manage your context usage efficiently."
+            )
+            if messages and messages[0]["role"] == "system":
+                messages[0]["content"] = original_system_prompt + status_msg
+
+            msg = get_llm_msg(
+                model_config["id"],
+                messages,
+                verbose=verbose,
+                model_alias=model_config.get("alias"),
+                usage_tracker=usage_tracker,
+            )
+            calls = extract_calls(msg, turn)
+            if not calls:
+                final_answer = strip_think_tags(msg.get("content", ""))
+                if is_markdown(final_answer):
+                    console = Console()
+                    console.print(Markdown(final_answer))
+                else:
+                    print(final_answer)
+                break
+            messages.append(msg)
+            for call in calls:
+                result = dispatch_tool_call(call, model_config["max_chars"], summarize)
+                messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": call["id"],
+                        "content": json.dumps(result),
+                    }
+                )
+        if turn >= MAX_TURNS:
+            print("Error: Max turns reached.")
+    except Exception as e:
+        print(f"Error: {str(e)}")
+    finally:
+        print(f"\nQuery completed in {time.perf_counter() - start_time:.2f} seconds")
+    return final_answer
+
+
+def generate_summaries(
+    query: str, answer: str, usage_tracker: Optional[UsageTracker] = None
+) -> tuple[str, str]:
+    """Generate summaries for query and answer using the summarization model."""
+    query_summary = ""
+    answer_summary = ""
+
+    # Generate Query Summary (if needed)
+    if len(query) > CONTINUE_QUERY_THRESHOLD:
+        try:
+            msgs = [
+                {
+                    "role": "system",
+                    "content": SUMMARIZE_QUERY_PROMPT_TEMPLATE.format(
+                        QUERY_SUMMARY_MAX_CHARS=QUERY_SUMMARY_MAX_CHARS
+                    ),
+                },
+                {"role": "user", "content": query[:1000]},
+            ]
+            model_id = MODELS[SUMMARIZATION_MODEL]["id"]
+            model_alias = MODELS[SUMMARIZATION_MODEL].get("alias", SUMMARIZATION_MODEL)
+            msg = get_llm_msg(
+                model_id,
+                msgs,
+                tools=None,
+                model_alias=model_alias,
+                usage_tracker=usage_tracker,
+            )
+            query_summary = strip_think_tags(msg.get("content", "")).strip()
+            if len(query_summary) > QUERY_SUMMARY_MAX_CHARS:
+                query_summary = query_summary[: QUERY_SUMMARY_MAX_CHARS - 3] + "..."
+        except Exception as e:
+            print(f"Error summarizing query: {e}")
+            query_summary = query[:QUERY_SUMMARY_MAX_CHARS]
+
+    # Generate Answer Summary (Always)
+    try:
+        msgs = [
+            {
+                "role": "system",
+                "content": SUMMARIZE_ANSWER_PROMPT_TEMPLATE.format(
+                    ANSWER_SUMMARY_MAX_CHARS=ANSWER_SUMMARY_MAX_CHARS
+                ),
+            },
+            {"role": "user", "content": answer[:5000]},
+        ]
+        model_id = MODELS[SUMMARIZATION_MODEL]["id"]
+        model_alias = MODELS[SUMMARIZATION_MODEL].get("alias", SUMMARIZATION_MODEL)
+        msg = get_llm_msg(
+            model_id,
+            msgs,
+            tools=None,
+            model_alias=model_alias,
+            usage_tracker=usage_tracker,
+        )
+        answer_summary = strip_think_tags(msg.get("content", "")).strip()
+        if len(answer_summary) > ANSWER_SUMMARY_MAX_CHARS:
+            answer_summary = answer_summary[: ANSWER_SUMMARY_MAX_CHARS - 3] + "..."
+    except Exception as e:
+        print(f"Error summarizing answer: {e}")
+        answer_summary = answer[:ANSWER_SUMMARY_MAX_CHARS]
+
+    return query_summary, answer_summary

asky/storage.py

@@ -0,0 +1,157 @@
+"""SQLite database functions for conversation history."""
+
+import os
+import sqlite3
+from datetime import datetime
+from typing import List, Optional
+
+from asky.config import DB_PATH, CONTINUE_QUERY_THRESHOLD
+
+
+def init_db() -> None:
+    """Initialize the SQLite database and create tables if they don't exist."""
+    os.makedirs(DB_PATH.parent, exist_ok=True)
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute("""
+        CREATE TABLE IF NOT EXISTS history (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            timestamp TEXT,
+            query TEXT,
+            query_summary TEXT,
+            answer TEXT,
+            answer_summary TEXT,
+            model TEXT
+        )
+    """)
+    conn.commit()
+    conn.close()
+
+
+def get_history(limit: int = 10) -> List[tuple]:
+    """Retrieve recent history entries."""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute(
+        "SELECT id, timestamp, query, query_summary, answer_summary, model FROM history ORDER BY id DESC LIMIT ?",
+        (limit,),
+    )
+    rows = c.fetchall()
+    conn.close()
+    return rows
+
+
+def get_db_record_count() -> int:
+    """Return the total number of entries in the history table."""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    c.execute("SELECT COUNT(*) FROM history")
+    count = c.fetchone()[0]
+    conn.close()
+    return count
+
+
+def get_interaction_context(ids: List[int], full: bool = False) -> str:
+    """Get context from previous interactions by their IDs."""
+    if not ids:
+        return ""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    placeholders = ",".join(["?"] * len(ids))
+    query_str = f"SELECT id, query, query_summary, answer, answer_summary FROM history WHERE id IN ({placeholders})"
+    c.execute(query_str, ids)
+    results = c.fetchall()
+    conn.close()
+
+    context_parts = []
+    for row in results:
+        rid, query, q_sum, answer, a_sum = row
+        # Use summary if available and original query is long enough
+        if q_sum and len(query) >= CONTINUE_QUERY_THRESHOLD:
+            q_text = q_sum
+        else:
+            q_text = query
+
+        a_text = answer if full else a_sum
+        context_parts.append(f"Query {rid}: {q_text}\nAnswer {rid}: {a_text}")
+
+    return "\n\n".join(context_parts)
+
+
+def cleanup_db(target: Optional[str], delete_all: bool = False) -> None:
+    """Delete history records by ID, range, list, or all."""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+
+    try:
+        if delete_all:
+            c.execute("DELETE FROM history")
+            c.execute("DELETE FROM sqlite_sequence WHERE name='history'")
+            print("dataset cleaned completely.")
+        elif target:
+            _delete_by_target(c, target)
+        else:
+            print(
+                "Error: No target specified for cleanup. Use --all or provide IDs/range."
+            )
+
+        conn.commit()
+    except Exception as e:
+        print(f"Error during cleanup: {e}")
+    finally:
+        conn.close()
+
+
+def _delete_by_target(cursor: sqlite3.Cursor, target: str) -> None:
+    """Helper to delete records by target specification."""
+    # Check for range (e.g., "1-5")
+    if "-" in target:
+        try:
+            start, end = map(int, target.split("-"))
+            if start > end:
+                start, end = end, start
+            cursor.execute(
+                "DELETE FROM history WHERE id >= ? AND id <= ?", (start, end)
+            )
+            print(f"deleted records from {start} to {end}.")
+        except ValueError:
+            print("Error: Invalid range format. Use 'start-end' (e.g., 1-5).")
+    # Check for comma-separated list (e.g., "1,3,5")
+    elif "," in target:
+        try:
+            ids = [int(x.strip()) for x in target.split(",")]
+            placeholders = ",".join(["?"] * len(ids))
+            cursor.execute(f"DELETE FROM history WHERE id IN ({placeholders})", ids)
+            print(f"deleted records: {', '.join(map(str, ids))}.")
+        except ValueError:
+            print("Error: Invalid list format. Use comma-separated integers.")
+    # Single ID
+    else:
+        try:
+            rid = int(target)
+            cursor.execute("DELETE FROM history WHERE id = ?", (rid,))
+            print(f"deleted record {rid}.")
+        except ValueError:
+            print("Error: Invalid ID format. Must be an integer.")
+
+
+def save_interaction(
+    query: str,
+    answer: str,
+    model: str,
+    query_summary: str = "",
+    answer_summary: str = "",
+) -> None:
+    """Save an interaction to the database."""
+    conn = sqlite3.connect(DB_PATH)
+    c = conn.cursor()
+    timestamp = datetime.now().isoformat()
+    c.execute(
+        """
+        INSERT INTO history (timestamp, query, query_summary, answer, answer_summary, model)
+        VALUES (?, ?, ?, ?, ?, ?)
+    """,
+        (timestamp, query, query_summary, answer, answer_summary, model),
+    )
+    conn.commit()
+    conn.close()