Implement simple AI chat

Author: robin-rpr

templates/pages/ai.html

@@ -14,6 +14,9 @@
             <button id="clearChat" class="clear__button">Clear Chat</button>
         </div>
         <div class="chat__log" id="chat"></div>
+        <div class="context__usage" id="contextUsage" style="display: none;">
+            <span>Context: </span><span id="contextPercent">0%</span>
+        </div>
         <form class="chat__form" id="form">
             <input class="form__input" id="input" placeholder="Ask anything about networks" autocomplete="off" />
             <button class="form__send" type="submit">Send</button>
@@ -64,6 +67,9 @@
     background: #fff3e0;
     padding: 8px;
     margin: 5px 0;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
     border-radius: 4px;
     font-family: monospace;
     font-size: 0.9em;
@@ -76,6 +82,52 @@
     margin: 5px 0;
     border-radius: 4px;
 }
+
+.log__assistant {
+    background: #f8f9fa;
+    padding: 8px;
+    margin: 5px 0;
+    border-radius: 4px;
+    border-left: 3px solid #007bff;
+}
+
+.log__user {
+    background: #e8f5e8;
+    padding: 8px;
+    margin: 5px 0;
+    border-radius: 4px;
+    border-left: 3px solid #28a745;
+}
+
+.context__usage {
+    text-align: right;
+    padding: 5px 10px;
+    font-size: 0.85em;
+    color: #666;
+    background: #f8f9fa;
+    border-top: 1px solid #dee2e6;
+}
+
+.chat__form {
+    display: flex;
+    gap: 10px;
+    padding: 10px;
+}
+
+.form__input {
+    flex: 1;
+    padding: 8px;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+}
+
+.form__send {
+    padding: 8px 16px;
+    background: #007bff;
+    color: white;
+    border: none;
+    border-radius: 4px;
+}
 </style>
 
 <script>
@@ -84,17 +136,29 @@
 const input = document.getElementById('input')
 const modelSelect = document.getElementById('modelSelect')
 const clearButton = document.getElementById('clearChat')
+const contextUsage = document.getElementById('contextUsage')
+const contextPercent = document.getElementById('contextPercent')
 
 let history = []
 let ws = null
+let currentModel = null
+
+function updateTokenUsage(tokens, model) {
+    currentModel = model
+    contextUsage.style.display = 'block'
+    let maxTokens = 5000
+    
+    const percentage = Math.min((tokens / maxTokens) * 100, 100)
+    contextPercent.textContent = `${Math.round(percentage)}%`
+}
 
 async function initialize() {
     try {
         const response = await fetch('/api/v1/ai/models')
         const data = await response.json()
 
         if (data.models) {
-            modelSelect.innerHTML = '<option value="">Select a model (default: OpenRouter default)</option>'
+            modelSelect.innerHTML = '<option value="">Auto</option>'
             data.models.forEach(model => {
                 const option = document.createElement('option')
                 option.value = model.id
@@ -132,7 +196,13 @@
 
 function handle(data) {
     switch (data.type) {
+        case 'token_usage':
+            updateTokenUsage(data.tokens, data.model)
+            break
         case 'tool_call':
+            if (window.currentAssistantEl) {
+                window.currentAssistantEl = null
+            }
             append('thinking', `Calling ${data.name}...`)
             break
         case 'tool_result':
@@ -142,16 +212,20 @@
             if (!window.currentAssistantEl) {
                 window.currentAssistantEl = document.createElement('div')
                 window.currentAssistantEl.className = 'log__assistant'
-                window.currentAssistantEl.textContent = ''
+                window.currentAssistantEl.innerHTML = ''
                 chat.appendChild(window.currentAssistantEl)
                 chat.scrollTop = chat.scrollHeight
             }
-            window.currentAssistantEl.textContent += data.delta
+            // Convert newlines to <br> tags and escape HTML
+            const escapedDelta = data.delta.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;').replace(/\n/g, '<br>')
+            window.currentAssistantEl.innerHTML += escapedDelta
             chat.scrollTop = chat.scrollHeight
             break
         case 'assistant_done':
             if (window.currentAssistantEl) {
-                history.push({ role: 'assistant', content: window.currentAssistantEl.textContent })
+                // Store the raw text content (without HTML) in history
+                const rawContent = window.currentAssistantEl.textContent
+                history.push({ role: 'assistant', content: rawContent })
                 window.currentAssistantEl = null
             }
             break
@@ -187,6 +261,7 @@
     history = []
     chat.innerHTML = ''
     window.currentAssistantEl = null
+    contextUsage.style.display = 'none'
 })
 
 initialize()

templates/pages/index.html

@@ -46,14 +46,14 @@ <h1 class="container__heading">
                 </div>
             </form>
             <div class="container__cards">
-                <a class="cards__item">
+                <a class="cards__item" href="/ai">
                     <div class="item__title">AI</div>
                     <div class="item__description">
                         Supercharge your productivity with an easy-to-use interface
                     </div>
                     <i class="fa-sharp fa-light fa-sparkles"></i>
                 </a>
-                <a class="cards__item">
+                <a class="cards__item" href="/telescope">
                     <div class="item__title">Telescope</div>
                     <div class="item__description">
                         Query all our real-time data across the global routing table

utils/ai.py

@@ -2,6 +2,7 @@
 from typing import List, Dict, Any, Generator
 from utils.database import get_db, close_db
 from flask import current_app as app
+from datetime import datetime, date
 from config import Config
 from openai import OpenAI
 import json
@@ -29,20 +30,27 @@ def _build_tools_schema() -> List[Dict[str, Any]]:
         })
     return out_tools
 
+class DateTimeEncoder(json.JSONEncoder):
+    """Custom JSON encoder that handles datetime objects."""
+    def default(self, obj):
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+        return super().default(obj)
+
 def _execute_tool(name: str, arguments: Dict[str, Any]) -> str:
     """Execute MCP tool and return JSON result."""
     db = None
     try:
         db = get_db()
-        rows = mcp_run_tool(name, arguments or {}, db)
-        return json.dumps({'rows': rows})
+        result = mcp_run_tool(name, arguments or {}, db)
+        return json.dumps(result, cls=DateTimeEncoder)
     except Exception as e:
         app.logger.error('Tool execution failed: %s', str(e), exc_info=True)
-        return json.dumps({'error': str(e)})
+        return json.dumps({'error': str(e)}, cls=DateTimeEncoder)
     finally:
         close_db(db)
 
-def chat(messages: List[Dict[str, str]], model: str = None) -> Generator[Dict[str, Any], None, None]:
+def chat(messages: List[Dict[str, str]], model: str = "openrouter/auto") -> Generator[Dict[str, Any], None, None]:
     """
     Stream chat with OpenAI models using function calling.
     Let the LLM handle its own thinking process.
@@ -71,8 +79,9 @@ def chat(messages: List[Dict[str, str]], model: str = None) -> Generator[Dict[st
                 tools=tools,
                 tool_choice='auto',
                 temperature=0.2,
-                max_tokens=2000,
-                stream=True
+                max_tokens=5000,
+                stream=True,
+                stream_options={"include_usage": True},
             )
 
             # Process streaming response
@@ -83,6 +92,10 @@ def chat(messages: List[Dict[str, str]], model: str = None) -> Generator[Dict[st
             for chunk in response:
                 if not chunk.choices:
                     continue
+                
+                # Handle usage data from streaming response
+                if hasattr(chunk, 'usage') and chunk.usage:
+                    yield {'type': 'token_usage', 'tokens': chunk.usage.total_tokens, 'model': model}
 
                 delta = chunk.choices[0].delta
 
@@ -142,12 +155,13 @@ def chat(messages: List[Dict[str, str]], model: str = None) -> Generator[Dict[st
                     yield {'type': 'tool_result', 'name': name, 'output': output}
 
                     # Add tool result to conversation
-                    conversation.append({
+                    tool_message = {
                         'role': 'tool',
                         'tool_call_id': tool_call['id'],
                         'name': name,
                         'content': output
-                    })
+                    }
+                    conversation.append(tool_message)
 
                 # Continue to next round for tool calling
                 continue