Spaces:
Running
Running
| # mcp_servers.py (FIXED: Schema Enforcement + Detailed Logging + Usage Tracking) | |
| import asyncio | |
| import json | |
| import re | |
| import google.generativeai as genai | |
| from anthropic import AsyncAnthropic | |
| from openai import AsyncOpenAI | |
| from typing import Dict, Optional, Tuple, List, Any | |
| import config | |
| from utils import load_prompt | |
| from personas import PERSONAS_DATA | |
| EVALUATION_PROMPT_TEMPLATE = load_prompt(config.PROMPT_FILES["evaluator"]) | |
| # --- DEFINING THE SCHEMA TO FORCE JUSTIFICATIONS --- | |
| EVALUATION_SCHEMA = { | |
| "type": "OBJECT", | |
| "properties": { | |
| "Novelty": { | |
| "type": "OBJECT", | |
| "properties": { | |
| "score": {"type": "INTEGER"}, | |
| "justification": {"type": "STRING"} | |
| }, | |
| "required": ["score", "justification"] | |
| }, | |
| "Usefulness_Feasibility": { | |
| "type": "OBJECT", | |
| "properties": { | |
| "score": {"type": "INTEGER"}, | |
| "justification": {"type": "STRING"} | |
| }, | |
| "required": ["score", "justification"] | |
| }, | |
| "Flexibility": { | |
| "type": "OBJECT", | |
| "properties": { | |
| "score": {"type": "INTEGER"}, | |
| "justification": {"type": "STRING"} | |
| }, | |
| "required": ["score", "justification"] | |
| }, | |
| "Elaboration": { | |
| "type": "OBJECT", | |
| "properties": { | |
| "score": {"type": "INTEGER"}, | |
| "justification": {"type": "STRING"} | |
| }, | |
| "required": ["score", "justification"] | |
| }, | |
| "Cultural_Appropriateness": { | |
| "type": "OBJECT", | |
| "properties": { | |
| "score": {"type": "INTEGER"}, | |
| "justification": {"type": "STRING"} | |
| }, | |
| "required": ["score", "justification"] | |
| } | |
| }, | |
| "required": ["Novelty", "Usefulness_Feasibility", "Flexibility", "Elaboration", "Cultural_Appropriateness"] | |
| } | |
| def extract_json(text: str) -> dict: | |
| """Robustly extracts JSON from text.""" | |
| try: | |
| clean_text = text.strip() | |
| if "```json" in clean_text: | |
| clean_text = clean_text.split("```json")[1].split("```")[0].strip() | |
| elif "```" in clean_text: | |
| clean_text = clean_text.split("```")[1].split("```")[0].strip() | |
| return json.loads(clean_text) | |
| except (json.JSONDecodeError, IndexError): | |
| try: | |
| match = re.search(r'(\{[\s\S]*\})', text) | |
| if match: | |
| return json.loads(match.group(1)) | |
| except: | |
| pass | |
| raise ValueError(f"Could not extract JSON from response: {text[:100]}...") | |
| class BusinessSolutionEvaluator: | |
| def __init__(self, gemini_client: Optional[genai.GenerativeModel]): | |
| if not gemini_client: | |
| raise ValueError("BusinessSolutionEvaluator requires a Google/Gemini client.") | |
| self.gemini_model = gemini_client | |
| if "ERROR:" in EVALUATION_PROMPT_TEMPLATE: | |
| raise FileNotFoundError(EVALUATION_PROMPT_TEMPLATE) | |
| async def evaluate(self, problem: str, solution_text: str) -> Tuple[dict, dict]: | |
| """Returns (evaluation_dict, usage_dict)""" | |
| print(f"Evaluating solution (live): {solution_text[:50]}...") | |
| base_prompt = EVALUATION_PROMPT_TEMPLATE.format(problem=problem, solution_text=solution_text) | |
| schema_instruction = """ | |
| [IMPORTANT SYSTEM INSTRUCTION] | |
| Ignore any previous examples of JSON formatting in this prompt. | |
| You MUST strictly follow the Output Schema provided below. | |
| For EACH of the 5 metrics (Novelty, Usefulness_Feasibility, etc.), you must provide an object with TWO fields: | |
| 1. "score": An integer from 1 to 5. | |
| 2. "justification": A specific sentence explaining why you gave that score. | |
| Do not output a list. Return a single JSON object describing the solution above. | |
| """ | |
| final_prompt = base_prompt + schema_instruction | |
| usage = {"model": "Gemini", "input": 0, "output": 0} | |
| try: | |
| response = await self.gemini_model.generate_content_async( | |
| final_prompt, | |
| generation_config=genai.types.GenerationConfig( | |
| response_mime_type="application/json", | |
| response_schema=EVALUATION_SCHEMA | |
| ) | |
| ) | |
| # Capture Usage | |
| if hasattr(response, "usage_metadata"): | |
| usage["input"] = response.usage_metadata.prompt_token_count | |
| usage["output"] = response.usage_metadata.candidates_token_count | |
| v_fitness = extract_json(response.text) | |
| if not isinstance(v_fitness, (dict, list)): | |
| raise ValueError(f"Judge returned invalid type: {type(v_fitness)}") | |
| print(f"Evaluation complete (live): {v_fitness}") | |
| return v_fitness, usage | |
| except Exception as e: | |
| print(f"ERROR: BusinessSolutionEvaluator failed: {e}") | |
| return { | |
| "Novelty": {"score": 1, "justification": f"Error: {str(e)}"}, | |
| "Usefulness_Feasibility": {"score": 1, "justification": f"Error: {str(e)}"}, | |
| "Flexibility": {"score": 1, "justification": f"Error: {str(e)}"}, | |
| "Elaboration": {"score": 1, "justification": f"Error: {str(e)}"}, | |
| "Cultural_Appropriateness": {"score": 1, "justification": f"Error: {str(e)}"} | |
| }, usage | |
| class AgentCalibrator: | |
| def __init__(self, api_clients: dict, evaluator: BusinessSolutionEvaluator): | |
| self.evaluator = evaluator | |
| self.api_clients = {name: client for name, client in api_clients.items() if client} | |
| self.sponsor_llms = list(self.api_clients.keys()) | |
| print(f"AgentCalibrator initialized with enabled clients: {self.sponsor_llms}") | |
| async def calibrate_team(self, problem: str) -> Tuple[Dict[str, Any], List[str], List[Dict[str, Any]], List[Dict[str, Any]]]: | |
| print(f"Running LIVE calibration test for specialist team on {self.sponsor_llms}...") | |
| error_log = [] | |
| detailed_results = [] | |
| all_usage_stats = [] # Collect all usage data here | |
| if not self.sponsor_llms: | |
| raise Exception("AgentCalibrator cannot run: No LLM clients are configured.") | |
| if len(self.sponsor_llms) == 1: | |
| default_llm = self.sponsor_llms[0] | |
| print("Only one LLM available. Skipping calibration.") | |
| plan = { | |
| "Plant": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Plant"], "llm": default_llm}, | |
| "Implementer": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Implementer"], "llm": default_llm}, | |
| "Monitor": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Monitor"], "llm": default_llm} | |
| } | |
| return plan, error_log, [], [] | |
| roles_to_test = { | |
| role: PERSONAS_DATA[key]["description"] | |
| for role, key in config.CALIBRATION_CONFIG["roles_to_test"].items() | |
| } | |
| test_problem = f"For the business problem '{problem}', generate a single, brief, one-paragraph concept-level solution." | |
| tasks = [] | |
| for role, persona in roles_to_test.items(): | |
| for llm_name in self.sponsor_llms: | |
| tasks.append(self.run_calibration_test(problem, role, llm_name, persona, test_problem)) | |
| results = await asyncio.gather(*tasks) | |
| detailed_results = results | |
| # Flatten results to extract usage | |
| for res in results: | |
| if "usage_gen" in res: all_usage_stats.append(res["usage_gen"]) | |
| if "usage_eval" in res: all_usage_stats.append(res["usage_eval"]) | |
| best_llms = {} | |
| role_metrics = config.CALIBRATION_CONFIG["role_metrics"] | |
| for role in roles_to_test.keys(): | |
| best_score = -1 | |
| best_llm = self.sponsor_llms[0] | |
| for res in results: | |
| if res["role"] == role: | |
| if res.get("error"): | |
| error_log.append(f"Calibration failed for {res['llm']} (as {role}): {res['error']}") | |
| continue | |
| metric = role_metrics[role] | |
| # Robust Dict Access | |
| raw_score_data = res.get("score", {}) | |
| if not isinstance(raw_score_data, (dict, list)): raw_score_data = {} | |
| if isinstance(raw_score_data, list): raw_score_data = raw_score_data[0] if len(raw_score_data) > 0 else {} | |
| metric_data = raw_score_data.get(metric, {}) | |
| if not isinstance(metric_data, (dict, list)): metric_data = {} | |
| if isinstance(metric_data, list): metric_data = metric_data[0] if len(metric_data) > 0 else {} | |
| score = metric_data.get("score", 0) | |
| if score > best_score: | |
| best_score = score | |
| best_llm = res["llm"] | |
| best_llms[role] = best_llm | |
| team_plan = { | |
| "Plant": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Plant"], "llm": best_llms["Plant"]}, | |
| "Implementer": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Implementer"], "llm": best_llms["Implementer"]}, | |
| "Monitor": {"persona": config.CALIBRATION_CONFIG["roles_to_test"]["Monitor"], "llm": best_llms["Monitor"]} | |
| } | |
| print(f"Calibration complete (live). Team plan: {team_plan}") | |
| return team_plan, error_log, detailed_results, all_usage_stats | |
| async def run_calibration_test(self, problem, role, llm_name, persona, test_problem): | |
| print(f"...Calibrating {role} on {llm_name}...") | |
| client = self.api_clients[llm_name] | |
| # 1. Generate Solution (and get usage) | |
| solution, gen_usage = await get_llm_response(llm_name, client, persona, test_problem) | |
| if "Error generating response" in solution: | |
| return {"role": role, "llm": llm_name, "error": solution, "output": solution, "usage_gen": gen_usage} | |
| # 2. Evaluate Solution (and get usage) | |
| score, eval_usage = await self.evaluator.evaluate(problem, solution) | |
| return { | |
| "role": role, | |
| "llm": llm_name, | |
| "score": score, | |
| "output": solution, | |
| "usage_gen": gen_usage, | |
| "usage_eval": eval_usage | |
| } | |
| # --- Unified API Call Function --- | |
| async def get_llm_response(client_name: str, client, system_prompt: str, user_prompt: str) -> Tuple[str, dict]: | |
| """Returns (text_response, usage_dict)""" | |
| usage = {"model": client_name, "input": 0, "output": 0} | |
| try: | |
| if client_name == "Gemini": | |
| model = client | |
| full_prompt = [ | |
| {'role': 'user', 'parts': [system_prompt]}, | |
| {'role': 'model', 'parts': ["Understood. I will act as this persona."]}, | |
| {'role': 'user', 'parts': [user_prompt]} | |
| ] | |
| response = await model.generate_content_async(full_prompt) | |
| # Capture Gemini Usage | |
| if hasattr(response, "usage_metadata"): | |
| usage["input"] = response.usage_metadata.prompt_token_count | |
| usage["output"] = response.usage_metadata.candidates_token_count | |
| return response.text, usage | |
| elif client_name == "Anthropic": | |
| response = await client.messages.create( | |
| model=config.MODELS["Anthropic"]["default"], | |
| max_tokens=8192, | |
| system=system_prompt, | |
| messages=[{"role": "user", "content": user_prompt}] | |
| ) | |
| # Capture Anthropic Usage | |
| if hasattr(response, "usage"): | |
| usage["input"] = response.usage.input_tokens | |
| usage["output"] = response.usage.output_tokens | |
| return response.content[0].text, usage | |
| elif client_name == "SambaNova": | |
| completion = await client.chat.completions.create( | |
| model=config.MODELS["SambaNova"]["default"], | |
| messages=[ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt} | |
| ] | |
| ) | |
| # Capture SambaNova Usage | |
| if hasattr(completion, "usage"): | |
| usage["input"] = completion.usage.prompt_tokens | |
| usage["output"] = completion.usage.completion_tokens | |
| return completion.choices[0].message.content, usage | |
| except Exception as e: | |
| error_message = f"Error generating response from {client_name}: {str(e)}" | |
| print(f"ERROR: API call to {client_name} failed: {e}") | |
| return error_message, usage |