Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
-from telegram import Update
-from telegram.ext import ApplicationBuilder, MessageHandler, filters, ContextTypes
+import json
+import math
+from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
+from telegram.ext import ApplicationBuilder, MessageHandler, filters, ContextTypes, CallbackQueryHandler
 from transformers import pipeline
 
 pipe = pipeline(
@@ -8,28 +10,162 @@ pipe = pipeline(
     model="Qwen/Qwen2-0.5B-Instruct"
 )
 
-def build_prompt(msg):
-    return f"<|im_start|>user\n{msg}<|im_end|>\n<|im_start|>assistant\n"
+# 🧠 muistia käyttäjille
+memory = {}
+
+# 🛠️ TOOLS
+
+def tool_calculate(expression: str):
+    try:
+        return str(eval(expression))
+    except:
+        return "Error"
+
+def tool_search(query: str):
+    return f"(fake search result for: {query})"
+
+TOOLS = {
+    "calculate": tool_calculate,
+    "search": tool_search
+}
+
+# 🧠 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"
+    )
+
+# 🤖 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
+        )
+
+        text = out[0]["generated_text"]
+        reply = text.split("<|assistant|>")[-1].strip()
+
+        # yritä parse JSON
+        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:
+            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
-    prompt = build_prompt(user_msg)
-
-    out = pipe(
-        prompt,
-        max_new_tokens=80,
-        temperature=0.6,
-        top_p=0.9,
-        do_sample=True
-    )
 
-    text = out[0]["generated_text"]
-    reply = text.split("<|im_start|>assistant")[-1].strip()
+    # ⌨️ 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"],
+            reply_markup=InlineKeyboardMarkup(keyboard)
+        )
+
+        memory[user_id].append(f"Assistant: {data['text']}")
+        return
+
+    reply = result["text"]
+
+    memory[user_id].append(f"Assistant: {reply}")
 
     await update.message.reply_text(reply)
 
+# 🔘 BUTTON HANDLER
+
+async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    query = update.callback_query
+    await query.answer()
+
+    # treat as new user input
+    fake_update = Update(
+        update.update_id,
+        message=query.message
+    )
+    fake_update.message.text = query.data
+
+    await handle(fake_update, context)
+
+# 🚀 START
+
 app = ApplicationBuilder().token(os.environ["TELEGRAM_TOKEN"]).build()
 
 app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle))
+app.add_handler(CallbackQueryHandler(button_handler))
 
 app.run_polling()