Update app.py

app.py

@@ -2,7 +2,9 @@ import os
 import json
 import ast
 import math
+import time
 import asyncio
+import threading
 from collections import defaultdict, deque
 
 import wikipedia
@@ -22,47 +24,50 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
 MAX_HISTORY = 12
 MAX_STEPS = 4
+MAX_NEW_TOKENS_JSON = 220
+
+torch.set_num_threads(max(1, os.cpu_count() // 2))
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype="auto",
-    device_map="auto"
+    torch_dtype=torch.float32,
+    device_map="cpu",
+    low_cpu_mem_usage=True
 )
+model.eval()
 
 memory = defaultdict(lambda: deque(maxlen=MAX_HISTORY))
 button_state = defaultdict(dict)
 
-SYSTEM_PROMPT = """
+PLANNER_SYSTEM_PROMPT = """
 You are an advanced Telegram AI agent.
 
 You must never expose internal tool calls, JSON planning, scratch work, or control tokens to the user.
 
-You have three possible response formats.
+You have exactly three response formats.
 
-1) Final answer:
+1) Final:
 {"type":"final","text":"your message to the user"}
 
 2) Buttons:
 {"type":"buttons","text":"question for the user","buttons":[{"id":"choice_1","label":"Yes"},{"id":"choice_2","label":"No"}]}
 
-3) Tool call:
+3) Tool:
 {"type":"tool","name":"wiki_search","arguments":{"query":"Finland"}}
-
-or
-
 {"type":"tool","name":"calculate","arguments":{"expression":"(25*17)/5"}}
 
 Rules:
 - Output exactly one JSON object.
 - No markdown fences.
 - No extra text before or after JSON.
-- Use buttons when the user should choose between short options.
-- Use wiki_search for factual lookups, topics, places, people, concepts, summaries.
+- The button type must be exactly "buttons".
+- Use buttons when the user should choose between a few short options.
+- Use wiki_search for factual topics, places, people, concepts, summaries.
 - Use calculate for arithmetic or formula evaluation.
-- After receiving a tool result, continue and produce either another tool call, buttons, or a final answer.
+- If you are unsure, do not call a tool. Respond with a final answer instead.
+- After receiving a tool result, continue and respond with exactly one JSON object.
 - Prefer Finnish if the user speaks Finnish.
-- Keep answers clean and natural.
 - Never output tokens like <|end|>, <|im_start|>, <|im_end|>.
 """
 
@@ -103,7 +108,10 @@ def safe_calculate(expression: str) -> str:
         ast.List,
     )
 
-    tree = ast.parse(expression, mode="eval")
+    try:
+        tree = ast.parse(expression, mode="eval")
+    except Exception:
+        return "Virhe: laskua ei voitu lukea."
 
     for node in ast.walk(tree):
         if not isinstance(node, allowed_nodes):
@@ -144,54 +152,62 @@ def wiki_search(query: str) -> str:
                 page = wikipedia.page(query, auto_suggest=True)
                 return page.summary[:1200]
         except Exception:
-            return f"En löytänyt Wikipedia-tulosta haulle: {query}"
+            return f"En löytänyt hakutulosta haulle: {query}"
 
 TOOLS = {
     "wiki_search": wiki_search,
     "calculate": safe_calculate,
 }
 
+def clean_text(text: str) -> str:
+    bad_tokens = [
+        "<|end|>",
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|assistant|>",
+        "<|user|>",
+        "<|system|>",
+        "</s>",
+    ]
+    for token in bad_tokens:
+        text = text.replace(token, "")
+    return text.strip()
+
 def extract_json(text: str):
     text = text.strip()
     decoder = json.JSONDecoder()
-
     for i, ch in enumerate(text):
         if ch == "{":
             try:
                 obj, end = decoder.raw_decode(text[i:])
                 trailing = text[i + end:].strip()
                 if trailing:
-                    return None
+                    continue
                 return obj
             except Exception:
                 continue
     return None
 
-def clean_text(text: str) -> str:
-    bad = ["<|end|>", "<|im_start|>", "<|im_end|>", "<|assistant|>", "<|user|>", "<|system|>"]
-    for token in bad:
-        text = text.replace(token, "")
-    return text.strip()
-
-def build_messages(user_id: int, user_text: str):
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+def build_planner_messages(user_id: int, user_text: str):
+    messages = [{"role": "system", "content": PLANNER_SYSTEM_PROMPT}]
     for item in memory[user_id]:
         messages.append(item)
     messages.append({"role": "user", "content": user_text})
     return messages
 
-def generate_json_response(messages):
+def generate_chat_text(messages, max_new_tokens=220):
     prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=220,
+            max_new_tokens=max_new_tokens,
             do_sample=False,
+            use_cache=True,
             pad_token_id=tokenizer.eos_token_id
         )
     new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
@@ -199,14 +215,19 @@ def generate_json_response(messages):
     return clean_text(text)
 
 def run_agent(user_id: int, user_text: str):
-    messages = build_messages(user_id, user_text)
+    messages = build_planner_messages(user_id, user_text)
 
     for _ in range(MAX_STEPS):
-        raw = generate_json_response(messages)
+        raw = generate_chat_text(messages, max_new_tokens=MAX_NEW_TOKENS_JSON)
         data = extract_json(raw)
 
         if not isinstance(data, dict):
-            return {"type": "final", "text": clean_text(raw)}
+            messages.append({"role": "assistant", "content": raw})
+            messages.append({
+                "role": "user",
+                "content": 'Your previous response was invalid. Output ONLY one valid JSON object.'
+            })
+            continue
 
         response_type = data.get("type")
 
@@ -220,15 +241,21 @@ def run_agent(user_id: int, user_text: str):
             text = clean_text(str(data.get("text", "")))
             buttons = data.get("buttons", [])
             normalized = []
-            for b in buttons[:6]:
-                if isinstance(b, dict):
-                    bid = str(b.get("id", "")).strip()
-                    label = str(b.get("label", "")).strip()
-                    if bid and label:
-                        normalized.append({"id": bid[:32], "label": label[:40]})
+            if isinstance(buttons, list):
+                for b in buttons[:6]:
+                    if isinstance(b, dict):
+                        bid = str(b.get("id", "")).strip()
+                        label = str(b.get("label", "")).strip()
+                        if bid and label:
+                            normalized.append({"id": bid[:32], "label": label[:40]})
             if text and normalized:
                 return {"type": "buttons", "text": text, "buttons": normalized}
-            return {"type": "final", "text": "Valintoja ei voitu muodostaa oikein."}
+            messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
+            messages.append({
+                "role": "user",
+                "content": 'That buttons response was invalid. Output ONLY one valid JSON object.'
+            })
+            continue
 
         if response_type == "tool":
             name = data.get("name")
@@ -256,7 +283,11 @@ def run_agent(user_id: int, user_text: str):
             })
             continue
 
-        return {"type": "final", "text": "Sain epäkelvon agenttivastauksen."}
+        messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
+        messages.append({
+            "role": "user",
+            "content": 'That response type was invalid. Output ONLY one valid JSON object.'
+        })
 
     return {"type": "final", "text": "Pyyntö vaati liikaa välivaiheita."}
 
@@ -267,13 +298,62 @@ async def typing_loop(chat, stop_event: asyncio.Event):
         except Exception:
             return
         try:
-            await asyncio.wait_for(stop_event.wait(), timeout=4.0)
+            await asyncio.wait_for(stop_event.wait(), timeout=2.0)
         except asyncio.TimeoutError:
             pass
 
-async def process_user_text(update_or_query_message, context: ContextTypes.DEFAULT_TYPE, user_id: int, text: str):
+def chunk_text_for_stream(text: str):
+    words = text.split()
+    if not words:
+        return [""]
+
+    chunks = []
+    current = ""
+
+    for word in words:
+        candidate = f"{current} {word}".strip()
+        if len(candidate) >= 35:
+            chunks.append(candidate)
+            current = ""
+        else:
+            current = candidate
+
+    if current:
+        chunks.append(current)
+
+    return chunks
+
+async def stream_text_reply(message, text: str):
+    text = clean_text(text)
+    if not text:
+        text = " "
+
+    chunks = chunk_text_for_stream(text)
+    sent = await message.reply_text("...")
+    assembled = ""
+    last_edit_time = 0.0
+
+    for i, chunk in enumerate(chunks):
+        assembled = f"{assembled} {chunk}".strip()
+        now = time.time()
+
+        if i < len(chunks) - 1:
+            if now - last_edit_time < 0.55:
+                await asyncio.sleep(0.55 - (now - last_edit_time))
+
+        safe_text = assembled[:4096]
+        try:
+            await sent.edit_text(safe_text)
+            last_edit_time = time.time()
+        except Exception:
+            pass
+
+    return sent
+
+async def process_user_text(message, context: ContextTypes.DEFAULT_TYPE, user_id: int, text: str):
     stop_event = asyncio.Event()
-    typing_task = asyncio.create_task(typing_loop(update_or_query_message.chat, stop_event))
+    typing_task = asyncio.create_task(typing_loop(message.chat, stop_event))
+
     try:
         result = await asyncio.to_thread(run_agent, user_id, text)
     finally:
@@ -285,6 +365,7 @@ async def process_user_text(update_or_query_message, context: ContextTypes.DEFAU
     if result["type"] == "buttons":
         keyboard = []
         button_state[user_id] = {}
+
         for b in result["buttons"]:
             button_state[user_id][b["id"]] = b["label"]
             keyboard.append([InlineKeyboardButton(b["label"], callback_data=f"btn:{b['id']}")])
@@ -297,22 +378,28 @@ async def process_user_text(update_or_query_message, context: ContextTypes.DEFAU
             )
         })
 
-        await update_or_query_message.reply_text(
+        await message.reply_text(
             result["text"],
             reply_markup=InlineKeyboardMarkup(keyboard)
         )
         return
 
     reply = clean_text(result["text"])
-    memory[user_id].append({"role": "assistant", "content": json.dumps({"type": "final", "text": reply}, ensure_ascii=False)})
-    await update_or_query_message.reply_text(reply)
+    memory[user_id].append({
+        "role": "assistant",
+        "content": json.dumps({"type": "final", "text": reply}, ensure_ascii=False)
+    })
+
+    await stream_text_reply(message, reply)
 
 async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
     if not update.message or not update.effective_user:
         return
+
     text = (update.message.text or "").strip()
     if not text:
         return
+
     await process_user_text(update.message, context, update.effective_user.id, text)
 
 async def handle_button(update: Update, context: ContextTypes.DEFAULT_TYPE):
@@ -339,7 +426,7 @@ async def handle_button(update: Update, context: ContextTypes.DEFAULT_TYPE):
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
     if not update.message:
         return
-    await update.message.reply_text("Moi. Olen AI-agentti. Voit kysyä mitä vain.")
+    await update.message.reply_text("Moi. Olen AI-agentti. Laita viestiä.")
 
 async def reset_chat(update: Update, context: ContextTypes.DEFAULT_TYPE):
     if not update.message or not update.effective_user:
@@ -351,6 +438,7 @@ async def reset_chat(update: Update, context: ContextTypes.DEFAULT_TYPE):
 
 def main():
     token = os.environ["TELEGRAM_TOKEN"]
+
     app = ApplicationBuilder().token(token).build()
 
     app.add_handler(CommandHandler("start", start))