tooluniverse 1.0.5__py3-none-any.whl → 1.0.7__py3-none-any.whl

tooluniverse/memory_manager.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""
+Memory Manager for Multi-Agent Systems
+Manages session-based memory for multi-agent workflows
+"""
+
+import json
+import uuid
+from datetime import datetime
+from typing import Dict, Any, Optional, List
+import threading
+
+
+class MemoryManager:
+    """专门管理multi-agent系统的记忆"""
+
+    def __init__(self, max_sessions=1000, session_timeout=3600):
+        self.sessions = {}  # 存储所有会话
+        self.max_sessions = max_sessions
+        self.session_timeout = session_timeout  # 会话超时时间（秒）
+        self.lock = threading.Lock()  # 线程安全
+
+    def create_session(self, user_id: str = None, session_name: str = None) -> str:
+        """创建新的会话"""
+        with self.lock:
+            session_id = f"{user_id or 'anonymous'}_{uuid.uuid4().hex[:8]}_{int(datetime.now().timestamp())}"
+
+            self.sessions[session_id] = {
+                "session_id": session_id,
+                "user_id": user_id or "anonymous",
+                "session_name": session_name
+                or f"Session_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
+                "created_at": datetime.now(),
+                "last_accessed": datetime.now(),
+                "context": {
+                    "user_query": "",
+                    "results": {},
+                    "history": [],
+                    "current_phase": "initializing",
+                },
+                "status": "active",
+            }
+
+            # 清理过期会话
+            self._cleanup_expired_sessions()
+
+            return session_id
+
+    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """获取会话信息"""
+        with self.lock:
+            if session_id in self.sessions:
+                self.sessions[session_id]["last_accessed"] = datetime.now()
+                return self.sessions[session_id]
+            return None
+
+    def update_session_context(self, session_id: str, updates: Dict[str, Any]) -> bool:
+        """更新会话上下文"""
+        with self.lock:
+            if session_id in self.sessions:
+                session = self.sessions[session_id]
+                session["context"].update(updates)
+                session["last_accessed"] = datetime.now()
+                return True
+            return False
+
+    def add_agent_result(
+        self, session_id: str, agent_name: str, result: Any, phase: str = None
+    ) -> bool:
+        """添加agent执行结果"""
+        with self.lock:
+            if session_id in self.sessions:
+                session = self.sessions[session_id]
+
+                # 添加到历史记录
+                history_entry = {
+                    "agent": agent_name,
+                    "result": result,
+                    "timestamp": datetime.now().isoformat(),
+                    "phase": phase or session["context"]["current_phase"],
+                }
+                session["context"]["history"].append(history_entry)
+
+                # 更新结果
+                session["context"]["results"][agent_name] = result
+                session["last_accessed"] = datetime.now()
+
+                return True
+            return False
+
+    def get_context_for_agent(self, session_id: str, agent_name: str) -> str:
+        """获取适合传递给agent的context信息"""
+        session = self.get_session(session_id)
+        if not session:
+            return "{}"
+
+        context = session["context"]
+
+        # 构建适合agent的context
+        agent_context = {
+            "session_id": session_id,
+            "user_id": session["user_id"],
+            "session_name": session["session_name"],
+            "user_query": context["user_query"],
+            "current_phase": context["current_phase"],
+            "previous_results": context["results"],
+            "execution_history": context["history"][-5:],  # 只保留最近5条历史
+            "timestamp": datetime.now().isoformat(),
+        }
+
+        return json.dumps(agent_context, indent=2)
+
+    def set_current_phase(self, session_id: str, phase: str) -> bool:
+        """设置当前执行阶段"""
+        return self.update_session_context(session_id, {"current_phase": phase})
+
+    def get_session_summary(self, session_id: str) -> Dict[str, Any]:
+        """获取会话摘要"""
+        session = self.get_session(session_id)
+        if not session:
+            return {}
+
+        return {
+            "session_id": session_id,
+            "user_id": session["user_id"],
+            "session_name": session["session_name"],
+            "created_at": session["created_at"].isoformat(),
+            "last_accessed": session["last_accessed"].isoformat(),
+            "current_phase": session["context"]["current_phase"],
+            "agents_executed": list(session["context"]["results"].keys()),
+            "history_count": len(session["context"]["history"]),
+            "status": session["status"],
+        }
+
+    def list_user_sessions(self, user_id: str) -> List[Dict[str, Any]]:
+        """获取用户的所有会话"""
+        with self.lock:
+            user_sessions = []
+            for session_id, session in self.sessions.items():
+                if session["user_id"] == user_id:
+                    user_sessions.append(self.get_session_summary(session_id))
+            return sorted(user_sessions, key=lambda x: x["last_accessed"], reverse=True)
+
+    def _cleanup_expired_sessions(self):
+        """清理过期会话"""
+        current_time = datetime.now()
+        expired_sessions = []
+
+        for session_id, session in self.sessions.items():
+            if (current_time - session["last_accessed"]).seconds > self.session_timeout:
+                expired_sessions.append(session_id)
+
+        for session_id in expired_sessions:
+            del self.sessions[session_id]
+
+    def close_session(self, session_id: str) -> bool:
+        """关闭会话"""
+        with self.lock:
+            if session_id in self.sessions:
+                self.sessions[session_id]["status"] = "closed"
+                return True
+            return False
+
+
+# 全局记忆管理器实例
+memory_manager = MemoryManager()

tooluniverse/molecule_2d_tool.py

@@ -0,0 +1,274 @@
+"""
+Molecule 2D Visualization Tool
+==============================
+
+Tool for visualizing 2D molecular structures using RDKit.
+Supports SMILES, InChI, molecule names, and various output formats.
+"""
+
+import base64
+import requests
+import io
+from typing import Any, Dict, Optional
+from .visualization_tool import VisualizationTool
+from .tool_registry import register_tool
+
+
+@register_tool("Molecule2DTool")
+class Molecule2DTool(VisualizationTool):
+    """Tool for visualizing 2D molecular structures using RDKit."""
+
+    def run(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Generate 2D molecular structure visualization."""
+        try:
+            from rdkit import Chem
+            from rdkit.Chem import Draw
+            from rdkit.Chem import rdDepictor
+
+            # Extract parameters
+            smiles = arguments.get("smiles")
+            inchi = arguments.get("inchi")
+            molecule_name = arguments.get("molecule_name")
+            width = arguments.get("width", 400)
+            height = arguments.get("height", 400)
+            output_format = arguments.get("output_format", "png")
+            show_atom_numbers = arguments.get("show_atom_numbers", False)
+            show_bond_numbers = arguments.get("show_bond_numbers", False)
+
+            # Create molecule object
+            mol = None
+            input_type = ""
+            input_data = ""
+
+            if smiles:
+                mol = Chem.MolFromSmiles(smiles)
+                input_type = "SMILES"
+                input_data = smiles
+            elif inchi:
+                mol = Chem.MolFromInchi(inchi)
+                input_type = "InChI"
+                input_data = inchi
+            elif molecule_name:
+                # Try to resolve molecule name to SMILES using PubChem
+                smiles_resolved = self._resolve_molecule_name(molecule_name)
+                if smiles_resolved:
+                    mol = Chem.MolFromSmiles(smiles_resolved)
+                    input_type = "Molecule Name"
+                    input_data = f"{molecule_name} -> {smiles_resolved}"
+                else:
+                    return self.create_error_response(
+                        f"Could not resolve molecule name: {molecule_name}"
+                    )
+            else:
+                return self.create_error_response(
+                    "Either smiles, inchi, or molecule_name must be provided"
+                )
+
+            if mol is None:
+                return self.create_error_response(
+                    "Failed to create molecule from input"
+                )
+
+            # Generate 2D coordinates
+            rdDepictor.Compute2DCoords(mol)
+
+            # Generate image
+            if output_format.lower() == "svg":
+                img_data = Draw.MolToSVG(mol, size=(width, height))
+                static_image = base64.b64encode(img_data.encode("utf-8")).decode(
+                    "utf-8"
+                )
+            else:
+                # Generate PNG
+                img = Draw.MolToImage(mol, size=(width, height))
+
+                # Convert to base64
+                buffer = io.BytesIO()
+                img.save(buffer, format="PNG")
+                img_data = buffer.getvalue()
+                static_image = self.convert_to_base64_image(img_data, "PNG")
+
+            # Calculate molecular properties
+            mol_props = self._calculate_molecular_properties(mol)
+
+            # Create modern HTML content
+            html_content = self.create_molecule_2d_html(
+                static_image,
+                mol_props,
+                width,
+                height,
+                title=f"2D Molecular Structure: {input_data[:20]}{'...' if len(input_data) > 20 else ''}",
+            )
+
+            # Prepare metadata
+            metadata = {
+                "width": width,
+                "height": height,
+                "output_format": output_format,
+                "input_type": input_type,
+                "show_atom_numbers": show_atom_numbers,
+                "show_bond_numbers": show_bond_numbers,
+                "molecular_properties": mol_props,
+            }
+
+            return self.create_visualization_response(
+                html_content=html_content,
+                viz_type="molecule_2d",
+                data={
+                    "input_data": input_data,
+                    "molecular_properties": mol_props,
+                    "smiles": Chem.MolToSmiles(mol) if mol else None,
+                },
+                static_image=static_image,
+                metadata=metadata,
+            )
+
+        except ImportError:
+            return self.create_error_response(
+                "RDKit is not installed. Please install it with: " "pip install rdkit",
+                "MissingDependency",
+            )
+        except Exception as e:
+            return self.create_error_response(
+                f"Failed to create molecule visualization: {str(e)}"
+            )
+
+    def _resolve_molecule_name(self, name: str) -> Optional[str]:
+        """Resolve molecule name to SMILES using PubChem."""
+        try:
+            # Use PubChem PUG REST API
+            url = (
+                f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/"
+                f"{name}/property/IsomericSMILES/JSON"
+            )
+            response = requests.get(url, timeout=10)
+            if response.status_code == 200:
+                data = response.json()
+                if "PropertyTable" in data and "Properties" in data["PropertyTable"]:
+                    props = data["PropertyTable"]["Properties"]
+                    if props and "IsomericSMILES" in props[0]:
+                        return props[0]["IsomericSMILES"]
+        except Exception:
+            pass
+        return None
+
+    def _calculate_molecular_properties(self, mol) -> Dict[str, Any]:
+        """Calculate basic molecular properties."""
+        try:
+            from rdkit import Chem
+            from rdkit.Chem import rdMolDescriptors
+
+            return {
+                "molecular_weight": rdMolDescriptors.CalcExactMolWt(mol),
+                "logp": rdMolDescriptors.CalcCrippenDescriptors(mol)[0],
+                "hbd": rdMolDescriptors.CalcNumHBD(mol),
+                "hba": rdMolDescriptors.CalcNumHBA(mol),
+                "tpsa": rdMolDescriptors.CalcTPSA(mol),
+                "rotatable_bonds": rdMolDescriptors.CalcNumRotatableBonds(mol),
+                "aromatic_rings": rdMolDescriptors.CalcNumAromaticRings(mol),
+                "heavy_atoms": mol.GetNumHeavyAtoms(),
+                "formal_charge": Chem.rdmolops.GetFormalCharge(mol),
+            }
+        except Exception:
+            return {}
+
+    def _create_molecule_html(
+        self,
+        mol,
+        input_data: str,
+        input_type: str,
+        width: int,
+        height: int,
+        static_image: str,
+        output_format: str,
+    ) -> str:
+        """Create HTML content for molecule visualization."""
+        try:
+            from rdkit import Chem
+
+            smiles = Chem.MolToSmiles(mol)
+            mol_props = self._calculate_molecular_properties(mol)
+
+            # Create properties table
+            props_html = ""
+            if mol_props:
+                props_html = (
+                    "<table border='1' "
+                    "style='border-collapse: collapse; margin: 10px 0;'>"
+                )
+                props_html += "<tr><th>Property</th><th>Value</th></tr>"
+                for prop, value in mol_props.items():
+                    if isinstance(value, float):
+                        value = f"{value:.2f}"
+                    props_html += f"<tr><td>{prop}</td><td>{value}</td></tr>"
+                props_html += "</table>"
+
+            return f"""
+            <!DOCTYPE html>
+            <html>
+            <head>
+                <meta charset="utf-8">
+                <title>2D Molecule Visualization</title>
+                <style>
+                    body {{
+                        font-family: Arial, sans-serif;
+                        margin: 20px;
+                        max-width: 1000px;
+                    }}
+                    .molecule-container {{
+                        border: 1px solid #ccc;
+                        border-radius: 5px;
+                        padding: 20px;
+                        margin: 10px 0;
+                        text-align: center;
+                    }}
+                    .molecule-image {{
+                        margin: 10px 0;
+                    }}
+                    .properties {{
+                        margin: 20px 0;
+                        text-align: left;
+                    }}
+                    .info {{
+                        background-color: #f5f5f5;
+                        padding: 10px;
+                        border-radius: 5px;
+                        margin: 10px 0;
+                    }}
+                </style>
+            </head>
+            <body>
+                <h2>2D Molecular Structure Visualization</h2>
+
+                <div class="info">
+                    <h3>Input Information</h3>
+                    <p><strong>Type:</strong> {input_type}</p>
+                    <p><strong>Data:</strong> {input_data}</p>
+                    <p><strong>SMILES:</strong> {smiles}</p>
+                </div>
+
+                <div class="molecule-container">
+                    <h3>Molecular Structure</h3>
+                    <div class="molecule-image">
+                        <img src="data:image/{output_format.lower()};base64,{static_image}"
+                             alt="2D Molecular Structure"
+                             style="max-width: 100%; height: auto;" />
+                    </div>
+                </div>
+
+                <div class="properties">
+                    <h3>Molecular Properties</h3>
+                    {props_html}
+                </div>
+
+                <div class="info">
+                    <h3>Visualization Details</h3>
+                    <p><strong>Dimensions:</strong> {width} × {height} "
+                    "pixels</p>
+                    <p><strong>Format:</strong> {output_format.upper()}</p>
+                </div>
+            </body>
+            </html>
+            """
+        except Exception as e:
+            return f"<div class='error'>Error creating HTML: {str(e)}</div>"