real_ai_agents_implementation.py

#!/usr/bin/env python3
"""
True Agentic Implementation for Patent Architect AI v2
This version implements a genuine, stateful, multi-agent negotiation workflow
where agent outputs dynamically influence subsequent agent actions.
"""

import os
import json
import re
import base64
import requests
from typing import Dict, List, Optional, Tuple, Generator
from dataclasses import dataclass, field
import google.generativeai as genai
from dotenv import load_dotenv

# --- Tool Imports ---
def web_search(query: str, num_results: int = 5) -> List[Dict]:
    """Performs a real web search using the free Serper.dev API."""
    
    api_key = os.getenv("SERPER_API_KEY")

    if not api_key:
        print("⚠️ WARNING: SERPER_API_KEY not set. Web search is disabled.")
        return [{"title": "Web Search Disabled", "link": "#", "snippet": "Please provide a Serper.dev API key to enable live web search."}]

    url = "https://google.serper.dev/search"
    payload = json.dumps({
        "q": query,
        "num": num_results
    })
    headers = {
        'X-API-KEY': api_key,
        'Content-Type': 'application/json'
    }

    try:
        print(f"Executing REAL web search for: {query} via Serper.dev")
        response = requests.post(url, headers=headers, data=payload, timeout=10)
        response.raise_for_status()
        search_results = response.json()
        
        # The key for organic results is 'organic'
        if 'organic' not in search_results:
            return []

        # Format the results to match the expected structure
        formatted_results = [
            {
                "title": item.get('title'),
                "link": item.get('link'),
                "snippet": item.get('snippet')
            }
            for item in search_results.get('organic', [])
        ]
        return formatted_results

    except requests.exceptions.RequestException as e:
        print(f"Error during web search with Serper.dev: {e}")
        return [{"title": "Web Search Error", "link": "#", "snippet": f"An error occurred during the search: {e}"}]


# --- Configuration ---
load_dotenv()
# Configure Gemini
try:
    genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
    GEMINI_AVAILABLE = True
except (ValueError, TypeError) as e:
    print(f"Gemini API key not found or invalid: {e}")
    GEMINI_AVAILABLE = False


# --- Data Structures ---
@dataclass
class NegotiationState:
    invention_disclosure: str
    key_concepts: List[str] = field(default_factory=list)
    prior_art_analysis: Dict = field(default_factory=dict)
    strategic_mandate: str = ""
    technical_summary: str = ""
    patent_claims: str = ""
    figure_description: str = ""
    ideogram_image_b64: str = ""
    negotiation_transcript: List[Dict] = field(default_factory=list)

# --- Base Agent Class ---
class BaseAgent:
    def __init__(self, model_name='gemini-1.5-flash'):
        if not GEMINI_AVAILABLE:
            self.model = None
            return
        self.model = genai.GenerativeModel(model_name)

    def _execute_prompt(self, prompt: str) -> str:
        if not self.model:
            return f"Error: The '{self.__class__.__name__}' agent could not run because the Gemini API is not configured. Please set the GEMINI_API_KEY."
        try:
            response = self.model.generate_content(prompt)
            return response.text
        except Exception as e:
            print(f"Error executing prompt for {self.__class__.__name__}: {e}")
            return f"Error: Could not get a response from the model. Details: {e}"

# --- Specialized Agents ---
class PriorArtDetective(BaseAgent):
    def analyze(self, invention_disclosure: str) -> Dict:
        # Phase 1: Extract Key Concepts
        concept_prompt = f"""
        Analyze the following invention disclosure. Your task is to extract the most important, specific, and searchable technical keywords or concepts.
        - If the disclosure is detailed, extract up to 3-5 key concepts.
        - If the disclosure is short or a single phrase, the main phrase itself may be the best concept.
        
        Return these concepts as a simple JSON array of strings. Do not add any other text or explanation. Your output MUST be only the JSON array.

        Invention: "{invention_disclosure}"
        
        Example for detailed input:
        Invention: "My invention is a smart coffee mug that uses a novel phase-change material to keep coffee at a perfect temperature. It also has a mobile app that connects via Bluetooth to let the user set their preferred temperature. The key innovation is a machine learning algorithm that learns the user's drinking habits to pre-warm or cool the mug, optimizing energy use."
        Output:
        ["phase-change material thermal management", "predictive pre-warming algorithm for beverage container", "user habit machine learning for temperature control"]

        Example for short input:
        Invention: "Non-invasive glucose monitoring using Raman spectroscopy"
        Output:
        ["non-invasive glucose monitoring", "Raman spectroscopy for glucose detection"]
        """
        response_text = self._execute_prompt(concept_prompt)
        key_concepts = []
        try:
            # More robust JSON parsing
            match = re.search(r'\[(.*?)\]', response_text, re.DOTALL)
            if match:
                # Handle cases where the model might just return comma-separated strings without brackets
                key_concepts = json.loads(match.group(0))
            else:
                cleaned_response = response_text.strip().replace('"', '').replace("'", "")
                if cleaned_response:
                    key_concepts = [c.strip() for c in cleaned_response.split(',') if c.strip()]

        except json.JSONDecodeError:
            # The model likely failed to return JSON. Try to parse it as a simple list.
            cleaned_response = response_text.strip().replace('"', '').replace("'", "").replace('[', '').replace(']', '')
            if cleaned_response:
                key_concepts = [c.strip() for c in cleaned_response.split(',') if c.strip()]
        
        # Fallback for very short disclosures if LLM still fails
        if not key_concepts and len(invention_disclosure.split()) < 10:
            # Also correct common typos as a courtesy
            corrected_disclosure = invention_disclosure.replace("Non-invasice", "Non-invasive")
            key_concepts = [corrected_disclosure]


        if not key_concepts:
            return {"key_concepts": [], "real_prior_art": [], "landscape_summary": "Could not identify key concepts to search for."}

        # Phase 2: Perform Live Web Searches
        search_findings = []
        for concept in key_concepts:
            # Search Google Patents
            patent_query = f'"{concept}" site:patents.google.com'
            patent_results = web_search(patent_query)
            search_findings.extend(patent_results)
            
            # Search Wikipedia for general background
            wiki_query = f'"{concept}" site:en.wikipedia.org'
            wiki_results = web_search(wiki_query)
            search_findings.extend(wiki_results)

        # Phase 3: Summarize Real Findings
        if not search_findings:
            return {"key_concepts": key_concepts, "real_prior_art": [], "landscape_summary": "No relevant prior art found in web search."}

        summary_prompt = f"""
        You are a patent analyst. Below are raw web search results for an invention related to "{', '.join(key_concepts)}".
        Your job is to analyze these results and provide a concise summary.

        Search Results:
        {json.dumps(search_findings, indent=2)}

        Tasks:
        1.  Create a brief, 2-3 sentence "landscape_summary" assessing how crowded the field appears to be based on these results.
        2.  Extract the top 5 most relevant findings and list them in a "real_prior_art" array, including their "title" and "link".

        Return the response as a single JSON object with keys: "landscape_summary" and "real_prior_art".
        """
        summary_response_text = self._execute_prompt(summary_prompt)
        
        try:
            # Clean the response to handle markdown code blocks
            match = re.search(r'```json\s*([\s\S]*?)\s*```', summary_response_text)
            if match:
                json_str = match.group(1)
            else:
                json_str = summary_response_text
            
            summary_data = json.loads(json_str.strip())
            summary_data["key_concepts"] = key_concepts # Add concepts for context
            return summary_data
        except json.JSONDecodeError:
            print(f"Failed to parse JSON from LLM summary response: {summary_response_text}")
            return {
                "key_concepts": key_concepts,
                "real_prior_art": search_findings, # Return the raw findings
                "landscape_summary": "Error: The AI's analysis of the search results could not be parsed.",
            }

class ChiefStrategistAgent(BaseAgent):
    def formulate_strategy(self, invention_disclosure: str, prior_art_analysis: Dict) -> str:
        prompt = f"""
        You are a Chief Patent Strategist. Your job is to determine the strongest angle for a successful patent application.
        
        Invention Disclosure:
        "{invention_disclosure}"

        Prior Art Analysis:
        - Key Concepts: {prior_art_analysis.get('key_concepts', [])}
        - Real Prior Art Found: {prior_art_analysis.get('real_prior_art', [])}
        - Landscape Summary: {prior_art_analysis.get('landscape_summary', '')}

        Based on the REAL prior art found, formulate a clear, one-sentence "Strategic Mandate". This mandate must identify the single most patentable aspect of the invention that appears novel and non-obvious compared to the search results.

        Example Mandates:
        - "The strategic focus shall be on the novel method for data encryption, as the search results confirm existing hardware implementations."
        - "The patentability of this invention rests on the unique chemical composition of the coating, which appears distinct from the cited art."
        - "We will patent the specific algorithm for adaptive lighting control, as the general hardware is well-known according to the search."

        Formulate the Strategic Mandate for the provided invention.
        """
        return self._execute_prompt(prompt)

class TechnicalWriterAgent(BaseAgent):
    def write_summary(self, invention_disclosure: str, strategic_mandate: str) -> str:
        prompt = f"""
        You are a professional patent writer. Your task is to write a "Summary of the Invention" section for a patent application.
        
        Invention Disclosure: "{invention_disclosure}"
        
        **CRITICAL INSTRUCTION:** You must follow this Strategic Mandate provided by the Chief Strategist:
        **Strategic Mandate: "{strategic_mandate}"**

        Write a concise, professional summary (2-3 paragraphs). Ensure that the summary heavily emphasizes the aspect highlighted in the Strategic Mandate as the core of the invention.
        """
        return self._execute_prompt(prompt)

class ClaimsDrafterAgent(BaseAgent):
    def draft_claims(self, invention_disclosure: str, strategic_mandate: str) -> str:
        prompt = f"""
        You are a patent attorney specializing in claim drafting.
        
        Invention Disclosure: "{invention_disclosure}"
        
        **CRITICAL INSTRUCTION:** Your claim set MUST be aligned with the following mandate:
        **Strategic Mandate: "{strategic_mandate}"**

        Draft a set of 5-7 patent claims.
        - The independent claim (Claim 1) must be directly focused on the feature identified in the Strategic Mandate.
        - Dependent claims should add further specifics and variations.
        - Ensure the claims are clear, concise, and properly formatted.
        """
        return self._execute_prompt(prompt)
        
class FigureDrafterAgent(BaseAgent):
    def describe_figure(self, invention_disclosure: str, strategic_mandate: str) -> str:
        prompt = f"""
        You are a patent illustrator's assistant. You need to generate the LaTeX/TikZ code for a key technical figure.

        Invention Disclosure: "{invention_disclosure}"

        **CRITICAL INSTRUCTION:** The figure must visually represent the core idea from the mandate:
        **Strategic Mandate: "{strategic_mandate}"**

        1.  Decide on the best type of figure to illustrate the mandate (e.g., flowchart, system diagram, cross-section).
        2.  Generate the complete LaTeX/TikZ code to create this figure.

        **CRITICAL OUTPUT FORMAT:**
        Return ONLY the raw LaTeX code, starting with `\\documentclass` and ending with `\\end{document}`.
        DO NOT include any description, explanation, or any text outside of the LaTeX code block.
        """
        return self._execute_prompt(prompt)

class SegmindIdeogramAgent:
    def __init__(self):
        self.api_key = os.getenv("SEGMIND_API_KEY") # Use user's key as default
        self.url = "https://api.segmind.com/v1/ideogram-3"

    def generate_image(self, technical_summary: str, strategic_mandate: str) -> Optional[str]:
        if not self.api_key:
            return None

        # Create a more cinematic prompt for image generation
        image_prompt = f"""
        Create a photorealistic, cinematic photograph representing the following invention.
        The image should focus on the core concept defined by the strategic mandate.

        Invention Summary: "{technical_summary}"
        Core Concept (Strategic Mandate): "{strategic_mandate}"

        Translate this technical concept into a visually stunning and professional marketing image.
        Emphasize the most innovative aspect. For example, if it's an algorithm, show a sleek user interface or an abstract representation of data flow, not just the hardware.
        """
        
        data = {
          "prompt": image_prompt,
          "resolution": "1024x1024",
          "style_type": "REALISTIC"
        }
        headers = {'x-api-key': self.api_key}

        try:
            response = requests.post(self.url, json=data, headers=headers)
            if response.status_code == 200:
                return base64.b64encode(response.content).decode('utf-8')
            else:
                print(f"Segmind API Error: {response.status_code} - {response.text}")
                return None
        except Exception as e:
            print(f"Error calling Segmind API: {e}")
            return None


# --- The Orchestrator ---
class AgenticNegotiator:
    def __init__(self, invention_disclosure: str):
        self.state = NegotiationState(invention_disclosure=invention_disclosure)
        self.agents = {
            "Prior Art Detective": PriorArtDetective(),
            "Chief Strategist": ChiefStrategistAgent(),
            "Technical Writer": TechnicalWriterAgent(),
            "Claims Drafter": ClaimsDrafterAgent(),
            "Figure Drafter": FigureDrafterAgent(),
            "Conceptual Artist": SegmindIdeogramAgent(),
        }

    def _update_transcript(self, agent_name: str, message: str, data: Optional[Dict] = None):
        entry = {"agent": agent_name, "message": message, "data": data or {}}
        self.state.negotiation_transcript.append(entry)

    def run_negotiation(self) -> Generator[NegotiationState, None, None]:
        # Step 0: Check for Gemini API Key
        if not GEMINI_AVAILABLE:
            self._update_transcript("System", "CRITICAL ERROR: `GEMINI_API_KEY` is not set. The agentic workflow cannot proceed. Please configure the environment variable.")
            yield self.state
            return

        # Step 1: Prior Art Detective
        agent_name = "Prior Art Detective"
        self._update_transcript(agent_name, "Analyzing the invention to understand the technical landscape...")
        yield self.state
        
        prior_art_result = self.agents[agent_name].analyze(self.state.invention_disclosure)
        self.state.prior_art_analysis = prior_art_result
        self.state.key_concepts = prior_art_result.get("key_concepts", [])
        self._update_transcript(agent_name, f"Analysis complete. The landscape appears to be: {prior_art_result.get('landscape_summary', 'N/A')}", prior_art_result)
        yield self.state

        # Step 2: Chief Strategist
        agent_name = "Chief Strategist"
        self._update_transcript(agent_name, "Reviewing prior art to determine the most defensible patenting strategy...")
        yield self.state
        
        mandate = self.agents[agent_name].formulate_strategy(self.state.invention_disclosure, self.state.prior_art_analysis)
        self.state.strategic_mandate = mandate
        self._update_transcript(agent_name, f"Strategy formulated. All agents will now adhere to the following mandate: **{mandate}**")
        yield self.state

        # Step 3: Guided Content Generation
        # Technical Summary
        agent_name = "Technical Writer"
        self._update_transcript(agent_name, "Acknowledged. Drafting the technical summary to align with the strategic mandate.")
        yield self.state
        summary = self.agents[agent_name].write_summary(self.state.invention_disclosure, self.state.strategic_mandate)
        self.state.technical_summary = summary
        self._update_transcript(agent_name, "Technical summary drafted.")
        yield self.state
        
        # Patent Claims
        agent_name = "Claims Drafter"
        self._update_transcript(agent_name, "Understood. Drafting patent claims focused on the mandated novel aspect.")
        yield self.state
        claims = self.agents[agent_name].draft_claims(self.state.invention_disclosure, self.state.strategic_mandate)
        self.state.patent_claims = claims
        self._update_transcript(agent_name, "Patent claims drafted.")
        yield self.state

        # Figure Description (LaTeX)
        agent_name = "Figure Drafter"
        self._update_transcript(agent_name, "Affirmative. Designing a technical figure that visually represents the core strategic mandate.")
        yield self.state
        figure_desc = self.agents[agent_name].describe_figure(self.state.invention_disclosure, self.state.strategic_mandate)
        self.state.figure_description = figure_desc
        self._update_transcript(agent_name, "Technical figure description and LaTeX code generated.")
        yield self.state

        # Conceptual Image (Ideogram)
        agent_name = "Conceptual Artist"
        self._update_transcript(agent_name, "Now generating a high-fidelity conceptual image based on the strategy...")
        yield self.state
        image_b64 = self.agents[agent_name].generate_image(self.state.technical_summary, self.state.strategic_mandate)
        if image_b64:
            self.state.ideogram_image_b64 = image_b64
            self._update_transcript(agent_name, "Conceptual image generated successfully.")
        else:
            self._update_transcript(agent_name, "Failed to generate conceptual image. The API may be unavailable or the key may be invalid.")
        yield self.state


        self._update_transcript("AgenticNegotiator", "All tasks complete. The patent application is ready for assembly.")
        yield self.state

def test_agentic_negotiation():
    """Test the new agentic negotiation workflow."""
    print("🤖 Testing True Agentic Workflow")
    print("=" * 60)
    
    if not GEMINI_AVAILABLE:
        print("\n❌ Cannot run test: GEMINI_API_KEY is not configured.")
        return

    test_invention = """
    My invention is a smart coffee mug that uses a novel phase-change material to keep coffee at a perfect temperature. It also has a mobile app that connects via Bluetooth to let the user set their preferred temperature. The key innovation is a machine learning algorithm that learns the user's drinking habits to pre-warm or cool the mug, optimizing energy use.
    """
    
    negotiator = AgenticNegotiator(invention_disclosure=test_invention)
    
    final_state = None
    for i, state in enumerate(negotiator.run_negotiation()):
        print(f"\n--- Turn {i+1} ---")
        last_message = state.negotiation_transcript[-1]
        print(f"**{last_message['agent']}:** {last_message['message']}")
        final_state = state

    print("\n\n✅ Negotiation Complete!")
    print("=" * 60)
    print(f"\n**Final Strategic Mandate:**\n{final_state.strategic_mandate}")
    print(f"\n**Generated Claims Preview:**\n{final_state.patent_claims[:300]}...")
    print(f"\n**Generated Figure Description Preview:**\n{final_state.figure_description[:300]}...")
    if final_state.ideogram_image_b64:
        print(f"\n**Ideogram Image:** Generated successfully (Base64 data)")
    else:
        print(f"\n**Ideogram Image:** Failed to generate.")

if __name__ == "__main__":
    test_agentic_negotiation()