Update multi_agent.py

multi_agent.py

@@ -1,145 +1,86 @@
-"""
-Enhanced multi-agent system with dynamic marks configuration
-"""
-
 import json
-import time
-from typing import Dict, Any, Tuple
-from agents import GeneratorAgent, VerifierAgent, FormatterAgent, SearchAgent
-from prompts import get_generator_prompt, VERIFIER_PROMPT, FORMATTER_PROMPT
-
-class MultiAgentSystem:
-    """Enhanced orchestrator with dynamic marks support"""
-    
-    def __init__(self):
-        self.generator = GeneratorAgent()
-        self.verifier = VerifierAgent()
-        self.formatter = FormatterAgent()
-        self.searcher = SearchAgent()
-        
-        self.progress = {
-            "stage": "ready",
-            "message": "System ready",
-            "percentage": 0
-        }
-    
-    def update_progress(self, stage: str, message: str, percentage: int):
-        """Update generation progress"""
-        self.progress = {
-            "stage": stage,
-            "message": message,
-            "percentage": percentage
-        }
-        print(f"Progress: {stage} - {message} ({percentage}%)")
-    
-    def extract_syllabus_from_pdf(self, pdf_file) -> str:
-        """Extract text from uploaded syllabus PDF"""
+from agents import GroqClient, SerpClient
+from prompts import build_master_prompt
+
+
+GENERATOR_MODEL = "llama-3.1-70b-versatile"
+VERIFIER_MODEL = "gemma2-27b-it"
+FORMATTER_MODEL = "mixtral-8x7b-32768"
+
+
+class MultiAgentOrchestrator:
+    def __init__(self, groq: GroqClient, serp: SerpClient):
+        self.groq = groq
+        self.serp = serp
+
+    def run_pipeline(self, subject, stream, partA, partB, partC, syllabus_text, ref_qp_text):
+
+        # 1) Real-time search
+        serp_results = self.serp.search(f"{subject} latest 2024 2025")
+        snippets = "\n".join([r.get("snippet", "") for r in serp_results[:3]])
+
+        # 2) Build master prompt
+        master_prompt = build_master_prompt(
+            subject, stream, partA, partB, partC,
+            syllabus_text, ref_qp_text, snippets
+        )
+
+        # 3) AGENT 1 — GENERATOR
+        gen_output = self.groq.ask(
+            [{"role": "system", "content": "Exam generator agent."},
+             {"role": "user", "content": master_prompt}],
+            model=GENERATOR_MODEL
+        )
+
+        # Extract JSON from generator
         try:
-            import PyPDF2
-            pdf_reader = PyPDF2.PdfReader(pdf_file)
-            text = ""
-            for page in pdf_reader.pages:
-                text += page.extract_text()
-            return text[:3000]
-        except Exception as e:
-            print(f"PDF extraction error: {e}")
-            return "Syllabus text extraction failed. Using default template."
-    
-    def extract_reference_from_pdf(self, pdf_file) -> str:
-        """Extract text from reference question paper PDF"""
-        if pdf_file is None:
-            return "No reference question paper provided."
-        
+            json_start = gen_output.rfind("{")
+            qp_json = json.loads(gen_output[json_start:])
+        except:
+            qp_json = {"raw": gen_output}
+
+        # 4) AGENT 2 — VERIFIER
+        ver_prompt = f"Check this QP JSON for errors and propose corrections:\n{json.dumps(qp_json)}"
+        verifier_output = self.groq.ask(
+            [{"role": "system", "content": "Verifier agent."},
+             {"role": "user", "content": ver_prompt}],
+            model=VERIFIER_MODEL
+        )
+
         try:
-            import PyPDF2
-            pdf_reader = PyPDF2.PdfReader(pdf_file)
-            text = ""
-            for page in pdf_reader.pages[:3]:
-                text += page.extract_text()
-            return text[:2000]
-        except Exception as e:
-            print(f"Reference PDF extraction error: {e}")
-            return "Reference paper extraction failed."
-    
-    def generate_exam_package(self, 
-                            subject: str,
-                            stream: str,
-                            part_a_count: int,
-                            part_b_count: int, 
-                            part_c_count: int,
-                            part_a_marks: int = 2,
-                            part_b_marks: int = 13,
-                            part_c_marks: int = 14,
-                            syllabus_file = None,
-                            reference_file = None) -> Tuple[Dict[str, Any], str]:
-        """Enhanced main method with dynamic marks"""
-        
+            corrections = json.loads(verifier_output)
+        except:
+            corrections = {"notes": verifier_output}
+
+        # 5) AGENT 3 — FORMATTER
+        fmt_prompt = f"""
+Apply corrections to QP JSON and produce:
+1. final_qp
+2. answers
+3. obe
+All inside ONE valid JSON object.
+
+Original QP:
+{json.dumps(qp_json)}
+
+Corrections:
+{json.dumps(corrections)}
+        """
+
+        final_output = self.groq.ask(
+            [{"role": "system", "content": "Formatter agent."},
+             {"role": "user", "content": fmt_prompt}],
+            model=FORMATTER_MODEL
+        )
+
         try:
-            # Stage 1: Data Preparation
-            self.update_progress("preparation", "Extracting syllabus and reference data", 10)
-            syllabus_text = self.extract_syllabus_from_pdf(syllabus_file)
-            reference_text = self.extract_reference_from_pdf(reference_file)
-            
-            # Stage 2: Realtime Search
-            self.update_progress("search", "Fetching recent developments", 20)
-            realtime_updates = self.searcher.get_realtime_updates(subject)
-            
-            # Stage 3: Generation with dynamic marks
-            self.update_progress("generation", "Generating question paper structure", 40)
-            generator_prompt = get_generator_prompt(
-                subject=subject,
-                stream=stream,
-                syllabus_text=syllabus_text,
-                reference_text=reference_text,
-                realtime_updates=realtime_updates,
-                part_a_count=part_a_count,
-                part_b_count=part_b_count,
-                part_c_count=part_c_count,
-                part_a_marks=part_a_marks,
-                part_b_marks=part_b_marks,
-                part_c_marks=part_c_marks
-            )
-            
-            generated_content = self.generator.generate_question_paper(generator_prompt)
-            
-            if "error" in generated_content:
-                return {}, "Generation failed: " + generated_content["error"]
-            
-            # Stage 4: Verification with dynamic marks
-            self.update_progress("verification", "Verifying quality and standards", 60)
-            verifier_prompt = VERIFIER_PROMPT.format(
-                generated_content=json.dumps(generated_content, indent=2),
-                bloom_mix="60% R/U, 40% A/An/Ev" if stream == "CSE" else "50% R/U, 50% A/An/Ev",
-                part_a_count=part_a_count,
-                part_b_count=part_b_count, 
-                part_c_count=part_c_count,
-                part_a_marks=part_a_marks,
-                part_b_marks=part_b_marks,
-                part_c_marks=part_c_marks,
-                tag_requirements="Company tags" if stream == "CSE" else "GATE tags"
-            )
-            
-            verification_result = self.verifier.verify_content(verifier_prompt)
-            
-            # Stage 5: Formatting
-            self.update_progress("formatting", "Creating final outputs and answers", 80)
-            formatter_prompt = FORMATTER_PROMPT.format(
-                original_content=json.dumps(generated_content, indent=2),
-                corrections=json.dumps(verification_result, indent=2)
-            )
-            
-            final_output = self.formatter.format_final_output(formatter_prompt)
-            
-            # Stage 6: Completion
-            self.update_progress("completion", "Package generation complete", 100)
-            
-            return final_output, "Success"
-            
-        except Exception as e:
-            error_msg = f"System error: {str(e)}"
-            print(error_msg)
-            return {}, error_msg
-    
-    def get_progress(self) -> Dict[str, Any]:
-        """Get current progress"""
-        return self.progress
+            final_json = json.loads(final_output)
+        except:
+            final_json = {"raw": final_output}
+
+        return {
+            "generator_raw": gen_output,
+            "qp_json": qp_json,
+            "verifier": corrections,
+            "final": final_json
+        }