Update agent.py

agent.py

@@ -1,4 +1,4 @@
-# agent_v31.py
+# agent_v31.py (wersja generyczna – podejście uniwersalne bez ifów per pytanie)
 import os
 import re
 import io
@@ -24,159 +24,121 @@ class GaiaAgent:
         except Exception:
             return None, None
 
-    def get_web_info(self, query):
+    def search_web_context(self, question):
         try:
-            return self.search_tool.run(query)
+            return self.search_tool.run(question)
         except Exception:
             return "[NO WEB INFO FOUND]"
 
-    def ask(self, prompt, model="gpt-4-turbo"):
+    def ask(self, context, question, model="gpt-4-turbo"):
+        messages = [
+            {"role": "system", "content": "You are an expert assistant. Use provided web or file context to answer. Output only the short final answer, formatted correctly. Do not explain."},
+            {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}\n\nAnswer:"}
+        ]
         response = self.client.chat.completions.create(
             model=model,
-            messages=[
-                {"role": "system", "content": "Return only a short factual answer. Format it properly. Never guess."},
-                {"role": "user", "content": prompt.strip() + "\nAnswer:"}
-            ],
+            messages=messages,
             temperature=0.0,
         )
         return response.choices[0].message.content.strip()
 
-    def ask_audio(self, audio_bytes, question):
-        path = "/tmp/audio.mp3"
-        with open(path, "wb") as f:
-            f.write(audio_bytes)
-        transcript = self.client.audio.transcriptions.create(model="whisper-1", file=open(path, "rb"))
-        return self.ask(f"Audio transcript: {transcript.text}\n\n{question}")
-
-    def ask_image(self, image_bytes, question):
-        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
-        messages = [
-            {"role": "system", "content": "Return only the winning move in chess algebraic notation (e.g., Qd1). No explanation."},
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": question},
-                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}
-                ]
-            }
-        ]
-        response = self.client.chat.completions.create(model="gpt-4o", messages=messages)
-        return response.choices[0].message.content.strip()
-
-    def extract_from_excel(self, file_bytes):
-        try:
-            df = pd.read_excel(io.BytesIO(file_bytes), engine="openpyxl")
-            df.columns = [col.lower() for col in df.columns]
-            if 'category' in df.columns and 'sales' in df.columns:
-                df['sales'] = pd.to_numeric(df['sales'], errors='coerce')
-                food_df = df[df['category'].str.lower() == 'food']
-                total = food_df['sales'].sum()
-                return f"${total:.2f}" if not pd.isna(total) else "$0.00"
-        except Exception:
-            pass
-        return "$0.00"
-
-    def extract_commutative_set(self, question):
-        try:
-            rows = re.findall(r"\|([a-e])\|([a-e\|]+)\|", question)
-            table = {}
-            for row in rows:
-                key, values = row
-                table[key] = values.strip('|').split('|')
-            elements = list(table.keys())
-            non_comm = set()
-            for i, x in enumerate(elements):
-                for j, y in enumerate(elements):
-                    if x != y:
-                        a = table[x][j]
-                        b = table[y][i]
-                        if a != b:
-                            non_comm.update([x, y])
-            return ", ".join(sorted(non_comm))
-        except:
-            return ""
-
-    def extract_answer(self, raw, question):
+    def format_answer(self, answer, question):
         q = question.lower()
-        raw = raw.strip().strip("\"'").strip()
-
-        if "studio albums" in q:
-            try:
-                return str(w2n.word_to_num(raw))
-            except:
-                match = re.search(r"\b\d+\b", raw)
-                return match.group(0) if match else raw
-
-        if "algebraic notation" in q:
-            match = re.search(r"\b([KQBNR]?[a-h]?[1-8]?x?[a-h][1-8][+#]?)\b", raw)
-            return match.group(1) if match else raw
+        a = answer.strip().strip("\"'").strip()
 
         if "usd with two decimal places" in q:
-            match = re.search(r"\$?([0-9]+(?:\.[0-9]{1,2})?)", raw)
+            match = re.search(r"\$?([0-9]+(?:\.[0-9]{1,2})?)", a)
             return f"${float(match.group(1)):.2f}" if match else "$0.00"
 
+        if "algebraic notation" in q:
+            match = re.search(r"\b([KQBNR]?[a-h]?[1-8]?x?[a-h][1-8][+#]?)\b", a)
+            return match.group(1) if match else a
+
         if "ioc country code" in q:
-            match = re.search(r"\b[A-Z]{3}\b", raw.upper())
+            match = re.search(r"\b[A-Z]{3}\b", a.upper())
             return match.group(0)
 
+        if "first name" in q:
+            return a.split()[0]
+
         if "page numbers" in q:
-            pages = sorted(set(re.findall(r"\b\d+\b", raw)))
-            return ", ".join(pages)
+            nums = sorted(set(re.findall(r"\b\d+\b", a)))
+            return ", ".join(nums)
 
         if "at bats" in q:
-            match = re.search(r"\b(\d{3,4})\b", raw)
-            return match.group(1)
+            match = re.search(r"\b(\d{3,4})\b", a)
+            return match.group(1) if match else a
 
-        if "first name" in q:
-            return raw.split()[0]
+        if "studio albums" in q or "how many" in q:
+            try:
+                return str(w2n.word_to_num(a))
+            except:
+                match = re.search(r"\b\d+\b", a)
+                return match.group(0) if match else a
 
         if "award number" in q:
-            match = re.search(r"80NSSC[0-9A-Z]{6,7}", raw)
-            return match.group(0) if match else raw
+            match = re.search(r"80NSSC[0-9A-Z]{6,7}", a)
+            return match.group(0) if match else a
 
         if "vegetables" in q or "ingredients" in q:
-            stopwords = set(["pure", "extract", "granulated", "sugar", "juice", "vanilla", "ripe", "fresh", "whole", "bean", "pinch", "cups", "salt", "water"])
-            tokens = [t.lower() for t in re.findall(r"[a-zA-Z]+", raw)]
-            clean = [t for t in tokens if t not in stopwords and len(t) > 2]
-            return ", ".join(sorted(set(clean)))
+            tokens = [t.lower() for t in re.findall(r"[a-zA-Z]+", a)]
+            blacklist = {"extract", "juice", "pure", "vanilla", "sugar", "granulated", "fresh", "ripe", "pinch", "water", "whole", "cups", "salt"}
+            clean = sorted(set(t for t in tokens if t not in blacklist and len(t) > 2))
+            return ", ".join(clean)
+
+        return a
 
-        return raw
+    def handle_file_context(self, file_bytes, ctype, question):
+        if not file_bytes:
+            return ""
+        if "image" in ctype:
+            image_b64 = base64.b64encode(file_bytes).decode("utf-8")
+            messages = [
+                {"role": "system", "content": "You're a visual reasoning assistant. Answer the question based on the image. Output only the move in algebraic notation."},
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": question},
+                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}
+                    ]
+                }
+            ]
+            response = self.client.chat.completions.create(model="gpt-4o", messages=messages)
+            return response.choices[0].message.content.strip()
+        elif "audio" in ctype or question.endswith(".mp3"):
+            path = "/tmp/audio.mp3"
+            with open(path, "wb") as f:
+                f.write(file_bytes)
+            transcript = self.client.audio.transcriptions.create(model="whisper-1", file=open(path, "rb"))
+            return transcript.text
+        elif "excel" in ctype or question.endswith(".xlsx"):
+            try:
+                df = pd.read_excel(io.BytesIO(file_bytes), engine="openpyxl")
+                df.columns = [c.lower() for c in df.columns]
+                df['sales'] = pd.to_numeric(df['sales'], errors='coerce')
+                food_df = df[df['category'].str.lower() == 'food']
+                total = food_df['sales'].sum()
+                return f"${total:.2f}" if not pd.isna(total) else "$0.00"
+            except Exception:
+                return "[EXCEL ERROR]"
+        else:
+            try:
+                return file_bytes.decode("utf-8")[:3000]
+            except:
+                return ""
 
     def __call__(self, question, task_id=None):
         file_bytes, ctype = None, ""
         if task_id:
             file_bytes, ctype = self.fetch_file(task_id)
 
-        try:
-            if "youtube.com" in question:
-                video_id = re.search(r"v=([\w-]+)", question)
-                if video_id:
-                    summary = self.get_web_info(f"youtube video transcript {video_id.group(1)}")
-                    return self.ask(f"Transcript: {summary}\n\n{question}")
-
-            if "malko competition" in question.lower():
-                search = self.get_web_info("malko competition winner yugoslavia after 1977 site:wikipedia.org")
-                return self.ask(f"Using the search result:\n{search}\n\n{question}")
-
-            if "commutative" in question:
-                result = self.extract_commutative_set(question)
-                return result
-
-            if file_bytes and "image" in ctype:
-                raw = self.ask_image(file_bytes, question)
-            elif file_bytes and ("audio" in ctype or task_id.endswith(".mp3")):
-                raw = self.ask_audio(file_bytes, question)
-            elif file_bytes and ("excel" in ctype or task_id.endswith(".xlsx")):
-                return self.extract_from_excel(file_bytes)
-            elif file_bytes:
-                try:
-                    text = file_bytes.decode("utf-8")
-                    raw = self.ask(f"Text content:\n{text[:3000]}\n\n{question}")
-                except:
-                    raw = "[UNREADABLE FILE CONTENT]"
-            else:
-                raw = self.ask(question)
-        except Exception as e:
-            return f"[ERROR: {e}]"
-
-        return self.extract_answer(raw, question)
+        file_context = self.handle_file_context(file_bytes, ctype, question)
+        if file_context and not file_context.startswith("$"):
+            raw = self.ask(file_context, question)
+        elif file_context.startswith("$"):
+            return file_context  # Excel result
+        else:
+            web_context = self.search_web_context(question)
+            raw = self.ask(web_context, question)
+
+        return self.format_answer(raw, question)