Update app.py

app.py

@@ -1,171 +1,364 @@
 import os
 import json
+import ast
 import math
+import asyncio
+from collections import defaultdict, deque
+
+import wikipedia
+import torch
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
-from telegram.ext import ApplicationBuilder, MessageHandler, filters, ContextTypes, CallbackQueryHandler
-from transformers import pipeline
+from telegram.constants import ChatAction
+from telegram.ext import (
+    ApplicationBuilder,
+    MessageHandler,
+    CommandHandler,
+    CallbackQueryHandler,
+    ContextTypes,
+    filters,
+)
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL_NAME = "Qwen/Qwen2-1.5B-Instruct"
+MAX_HISTORY = 12
+MAX_STEPS = 4
 
-pipe = pipeline(
-    "text-generation",
-    model="Qwen/Qwen2-0.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype="auto",
+    device_map="auto"
 )
 
-# 🧠 muistia käyttäjille
-memory = {}
+memory = defaultdict(lambda: deque(maxlen=MAX_HISTORY))
+button_state = defaultdict(dict)
+
+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.
+
+1) Final answer:
+{"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:
+{"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.
+- Use calculate for arithmetic or formula evaluation.
+- After receiving a tool result, continue and produce either another tool call, buttons, or a final answer.
+- Prefer Finnish if the user speaks Finnish.
+- Keep answers clean and natural.
+- Never output tokens like <|end|>, <|im_start|>, <|im_end|>.
+"""
+
+def safe_calculate(expression: str) -> str:
+    allowed_names = {
+        "abs": abs,
+        "round": round,
+        "min": min,
+        "max": max,
+        "pow": pow,
+        "sqrt": math.sqrt,
+        "sin": math.sin,
+        "cos": math.cos,
+        "tan": math.tan,
+        "pi": math.pi,
+        "e": math.e,
+    }
+
+    allowed_nodes = (
+        ast.Expression,
+        ast.BinOp,
+        ast.UnaryOp,
+        ast.Num,
+        ast.Constant,
+        ast.Add,
+        ast.Sub,
+        ast.Mult,
+        ast.Div,
+        ast.FloorDiv,
+        ast.Mod,
+        ast.Pow,
+        ast.USub,
+        ast.UAdd,
+        ast.Load,
+        ast.Call,
+        ast.Name,
+        ast.Tuple,
+        ast.List,
+    )
+
+    tree = ast.parse(expression, mode="eval")
 
-# 🛠️ TOOLS
+    for node in ast.walk(tree):
+        if not isinstance(node, allowed_nodes):
+            return "Virhe: laskua ei voitu suorittaa turvallisesti."
+        if isinstance(node, ast.Call):
+            if not isinstance(node.func, ast.Name):
+                return "Virhe: laskua ei voitu suorittaa turvallisesti."
+            if node.func.id not in allowed_names:
+                return "Virhe: laskua ei voitu suorittaa turvallisesti."
+        if isinstance(node, ast.Name):
+            if node.id not in allowed_names:
+                return "Virhe: laskua ei voitu suorittaa turvallisesti."
 
-def tool_calculate(expression: str):
     try:
-        return str(eval(expression))
-    except:
-        return "Error"
+        result = eval(compile(tree, "<expr>", "eval"), {"__builtins__": {}}, allowed_names)
+        return str(result)
+    except Exception:
+        return "Virhe: laskua ei voitu suorittaa."
+
+def wiki_search(query: str) -> str:
+    query = query.strip()
+    if not query:
+        return "Virhe: tyhjä hakukysely."
 
-def tool_search(query: str):
-    return f"(fake search result for: {query})"
+    try:
+        wikipedia.set_lang("fi")
+        try:
+            return wikipedia.summary(query, sentences=3, auto_suggest=True)
+        except Exception:
+            page = wikipedia.page(query, auto_suggest=True)
+            return page.summary[:1200]
+    except Exception:
+        try:
+            wikipedia.set_lang("en")
+            try:
+                return wikipedia.summary(query, sentences=3, auto_suggest=True)
+            except Exception:
+                page = wikipedia.page(query, auto_suggest=True)
+                return page.summary[:1200]
+        except Exception:
+            return f"En löytänyt Wikipedia-tulosta haulle: {query}"
 
 TOOLS = {
-    "calculate": tool_calculate,
-    "search": tool_search
+    "wiki_search": wiki_search,
+    "calculate": safe_calculate,
 }
 
-# 🧠 PROMPT
-
-def build_prompt(user_id, msg):
-    history = memory.get(user_id, [])
-    history_text = "\n".join(history[-5:])
-
-    return (
-        "<|system|>\n"
-        "You are an AI agent.\n"
-        "You can:\n"
-        "- Answer normally\n"
-        "- Use tools\n"
-        "- Create buttons\n\n"
-
-        "TOOLS FORMAT:\n"
-        '{"type":"tool","name":"calculate","input":"2+2"}\n'
-        '{"type":"tool","name":"search","input":"weather Helsinki"}\n\n'
-
-        "BUTTON FORMAT:\n"
-        '{"type":"buttons","text":"Choose:","buttons":["A","B"]}\n\n'
-
-        "RULES:\n"
-        "- Only output JSON when using tools or buttons\n"
-        "- No explanation around JSON\n"
-        "- Otherwise reply normally\n"
-        "- Use tools when useful\n"
-        "- Use buttons for choices\n"
-        "<|end|>\n"
-
-        f"{history_text}\n"
-
-        "<|user|>\n"
-        f"{msg}"
-        "<|end|>\n"
-        "<|assistant|>\n"
+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
+                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}]
+    for item in memory[user_id]:
+        messages.append(item)
+    messages.append({"role": "user", "content": user_text})
+    return messages
+
+def generate_json_response(messages):
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
     )
-
-# 🤖 AGENT LOOP
-
-def run_agent(user_id, msg):
-    prompt = build_prompt(user_id, msg)
-
-    for _ in range(3):  # max 3 tool calls
-        out = pipe(
-            prompt,
-            max_new_tokens=120,
-            temperature=0.6,
-            top_p=0.9,
-            do_sample=True
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=220,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id
         )
-
-        text = out[0]["generated_text"]
-        reply = text.split("<|assistant|>")[-1].strip()
-
-        # yritä parse JSON
+    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    text = tokenizer.decode(new_tokens, skip_special_tokens=False)
+    return clean_text(text)
+
+def run_agent(user_id: int, user_text: str):
+    messages = build_messages(user_id, user_text)
+
+    for _ in range(MAX_STEPS):
+        raw = generate_json_response(messages)
+        data = extract_json(raw)
+
+        if not isinstance(data, dict):
+            return {"type": "final", "text": clean_text(raw)}
+
+        response_type = data.get("type")
+
+        if response_type == "final":
+            text = clean_text(str(data.get("text", "")))
+            if not text:
+                text = "En saanut muodostettua vastausta."
+            return {"type": "final", "text": text}
+
+        if response_type == "buttons":
+            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 text and normalized:
+                return {"type": "buttons", "text": text, "buttons": normalized}
+            return {"type": "final", "text": "Valintoja ei voitu muodostaa oikein."}
+
+        if response_type == "tool":
+            name = data.get("name")
+            arguments = data.get("arguments", {})
+
+            if name not in TOOLS:
+                return {"type": "final", "text": f"Tuntematon työkalu: {name}"}
+
+            if not isinstance(arguments, dict):
+                return {"type": "final", "text": "Työkalun argumentit olivat virheelliset."}
+
+            if name == "wiki_search":
+                query = str(arguments.get("query", "")).strip()
+                result = TOOLS[name](query)
+            elif name == "calculate":
+                expression = str(arguments.get("expression", "")).strip()
+                result = TOOLS[name](expression)
+            else:
+                result = "Työkalua ei voitu suorittaa."
+
+            messages.append({"role": "assistant", "content": json.dumps(data, ensure_ascii=False)})
+            messages.append({
+                "role": "user",
+                "content": f"Tool result for {name}:\n{result}\nNow continue and respond with exactly one JSON object."
+            })
+            continue
+
+        return {"type": "final", "text": "Sain epäkelvon agenttivastauksen."}
+
+    return {"type": "final", "text": "Pyyntö vaati liikaa välivaiheita."}
+
+async def typing_loop(chat, stop_event: asyncio.Event):
+    while not stop_event.is_set():
         try:
-            data = json.loads(reply)
-
-            # 🔘 BUTTONS
-            if data.get("type") == "buttons":
-                return {"buttons": data}
-
-            # 🛠️ TOOL CALL
-            if data.get("type") == "tool":
-                name = data.get("name")
-                inp = data.get("input")
-
-                if name in TOOLS:
-                    result = TOOLS[name](inp)
-
-                    # lisää contextiin
-                    prompt += f"\nTool result: {result}\nContinue.\n"
-                    continue
-
-        except:
+            await chat.send_action(ChatAction.TYPING)
+        except Exception:
+            return
+        try:
+            await asyncio.wait_for(stop_event.wait(), timeout=4.0)
+        except asyncio.TimeoutError:
             pass
 
-        return {"text": reply}
-
-    return {"text": "I couldn't complete that."}
-
-# 📩 MESSAGE HANDLER
-
-async def handle(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    user_id = update.effective_user.id
-    user_msg = update.message.text
-
-    # ⌨️ typing
-    await update.message.chat.send_action("typing")
-
-    result = run_agent(user_id, user_msg)
-
-    # tallenna muisti
-    memory.setdefault(user_id, []).append(f"User: {user_msg}")
-
-    if "buttons" in result:
-        data = result["buttons"]
-
-        keyboard = [
-            [InlineKeyboardButton(b, callback_data=b)]
-            for b in data["buttons"]
-        ]
-
-        await update.message.reply_text(
-            data["text"],
+async def process_user_text(update_or_query_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))
+    try:
+        result = await asyncio.to_thread(run_agent, user_id, text)
+    finally:
+        stop_event.set()
+        await typing_task
+
+    memory[user_id].append({"role": "user", "content": text})
+
+    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']}")])
+
+        memory[user_id].append({
+            "role": "assistant",
+            "content": json.dumps(
+                {"type": "buttons", "text": result["text"], "buttons": result["buttons"]},
+                ensure_ascii=False
+            )
+        })
+
+        await update_or_query_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(f"Assistant: {data['text']}")
+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)
 
-    reply = result["text"]
+async def handle_button(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    query = update.callback_query
+    if not query or not update.effective_user:
+        return
 
-    memory[user_id].append(f"Assistant: {reply}")
+    await query.answer()
 
-    await update.message.reply_text(reply)
+    data = query.data or ""
+    if not data.startswith("btn:"):
+        return
 
-# 🔘 BUTTON HANDLER
+    button_id = data[4:]
+    user_id = update.effective_user.id
+    label = button_state[user_id].get(button_id)
 
-async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    query = update.callback_query
-    await query.answer()
+    if not label:
+        await query.message.reply_text("Tämä valinta ei ole enää voimassa.")
+        return
 
-    # treat as new user input
-    fake_update = Update(
-        update.update_id,
-        message=query.message
-    )
-    fake_update.message.text = query.data
+    await process_user_text(query.message, context, user_id, label)
+
+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 handle(fake_update, context)
+async def reset_chat(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    if not update.message or not update.effective_user:
+        return
+    user_id = update.effective_user.id
+    memory[user_id].clear()
+    button_state[user_id].clear()
+    await update.message.reply_text("Muisti nollattu.")
 
-# 🚀 START
+def main():
+    token = os.environ["TELEGRAM_TOKEN"]
+    app = ApplicationBuilder().token(token).build()
 
-app = ApplicationBuilder().token(os.environ["TELEGRAM_TOKEN"]).build()
+    app.add_handler(CommandHandler("start", start))
+    app.add_handler(CommandHandler("reset", reset_chat))
+    app.add_handler(CallbackQueryHandler(handle_button))
+    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
 
-app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle))
-app.add_handler(CallbackQueryHandler(button_handler))
+    app.run_polling(drop_pending_updates=True)
 
-app.run_polling()
+if __name__ == "__main__":
+    main()