devduck 0.1.0__py3-none-any.whl → 0.1.1766644714__py3-none-any.whl

devduck/tools/state_manager.py

@@ -0,0 +1,292 @@
+"""DevDuck State Manager - Time-travel for agent conversations"""
+
+import os
+import tempfile
+import dill
+from pathlib import Path
+from datetime import datetime
+from typing import Dict, Any
+from strands import tool
+
+base_dir = Path(os.getenv("DEVDUCK_HOME", tempfile.gettempdir()))
+states_dir = base_dir / ".devduck" / "states"
+states_dir.mkdir(parents=True, exist_ok=True)
+
+
+@tool
+def state_manager(
+    action: str,
+    state_file: str = None,
+    query: str = None,
+    metadata: dict = None,
+    agent=None,  # Parent agent injection
+) -> Dict[str, Any]:
+    """Agent state management with time-travel capabilities.
+
+    Inspired by cagataycali/research-agent state export pattern.
+
+    Actions:
+    - export: Save current agent state to pkl
+    - load: Load and display state from pkl
+    - list: List available saved states
+    - resume: Load state and continue with new query (ephemeral)
+    - modify: Update pkl file metadata
+    - delete: Remove saved state
+
+    Args:
+        action: Operation to perform
+        state_file: Path to pkl file (auto-generated for export)
+        query: New query for resume action
+        metadata: Additional metadata for export/modify
+        agent: Parent agent (auto-injected by Strands)
+
+    Returns:
+        Dict with status and content
+
+    Examples:
+        # Save current state
+        state_manager(action="export", metadata={"note": "before refactor"})
+
+        # List saved states
+        state_manager(action="list")
+
+        # Resume from previous state (ephemeral, no mutation)
+        state_manager(action="resume", state_file="~/.devduck/states/devduck_20250116_032000.pkl", query="continue analysis")
+
+        # Modify state metadata
+        state_manager(action="modify", state_file="path/to/state.pkl", metadata={"tags": ["important", "refactor"]})
+    """
+    try:
+        if action == "export":
+            # Capture current agent state
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            state_file = states_dir / f"devduck_{timestamp}.pkl"
+
+            # Safe state extraction (avoid complex objects)
+            state_data = {
+                "version": "1.0",
+                "timestamp": datetime.now().isoformat(),
+                "system_prompt": agent.system_prompt,
+                "tools": list(agent.tool_names),
+                "model": {
+                    "model_id": getattr(agent.model, "model_id", "unknown"),
+                    "temperature": getattr(agent.model, "temperature", None),
+                    "provider": getattr(agent.model, "provider", "unknown"),
+                },
+                "metadata": metadata or {},
+                "environment": {
+                    "cwd": str(Path.cwd()),
+                    "devduck_version": "0.6.0",
+                },
+            }
+
+            # Try to capture conversation history if available
+            if hasattr(agent, "conversation_history"):
+                state_data["conversation_history"] = agent.conversation_history
+            elif hasattr(agent, "messages"):
+                state_data["conversation_history"] = agent.messages
+
+            # Save with dill
+            with open(state_file, "wb") as f:
+                dill.dump(state_data, f)
+
+            size = state_file.stat().st_size
+            return {
+                "status": "success",
+                "content": [
+                    {
+                        "text": f"✅ State exported: {state_file}\n"
+                        f"📦 Size: {size} bytes\n"
+                        f"🔧 Tools: {len(state_data['tools'])}\n"
+                        f"📝 Metadata: {metadata or 'none'}"
+                    }
+                ],
+            }
+
+        elif action == "list":
+            # List all saved states
+            states = sorted(
+                states_dir.glob("devduck_*.pkl"),
+                key=lambda p: p.stat().st_mtime,
+                reverse=True,
+            )
+
+            if not states:
+                return {
+                    "status": "success",
+                    "content": [{"text": "No saved states found"}],
+                }
+
+            output = f"📚 Found {len(states)} saved states:\n\n"
+            for i, state_path in enumerate(states[:10], 1):  # Show last 10
+                try:
+                    with open(state_path, "rb") as f:
+                        state_data = dill.load(f)
+
+                    timestamp = state_data.get("timestamp", "unknown")
+                    tools_count = len(state_data.get("tools", []))
+                    meta = state_data.get("metadata", {})
+
+                    output += f"{i}. {state_path.name}\n"
+                    output += f"   📅 {timestamp}\n"
+                    output += f"   🔧 {tools_count} tools\n"
+                    if meta:
+                        output += f"   📝 {meta}\n"
+                    output += "\n"
+                except:
+                    output += f"{i}. {state_path.name} (corrupted)\n\n"
+
+            return {"status": "success", "content": [{"text": output}]}
+
+        elif action == "load":
+            # Load and display state
+            if not state_file:
+                return {
+                    "status": "error",
+                    "content": [{"text": "state_file required for load"}],
+                }
+
+            state_path = Path(state_file).expanduser()
+            if not state_path.exists():
+                return {
+                    "status": "error",
+                    "content": [{"text": f"State file not found: {state_path}"}],
+                }
+
+            with open(state_path, "rb") as f:
+                state_data = dill.load(f)
+
+            # Pretty format
+            output = f"📦 State: {state_path.name}\n\n"
+            output += f"📅 Timestamp: {state_data.get('timestamp')}\n"
+            output += f"🤖 Model: {state_data.get('model', {}).get('model_id')}\n"
+            output += f"🔧 Tools ({len(state_data.get('tools', []))}): {', '.join(state_data.get('tools', []))}\n"
+            output += f"📝 Metadata: {state_data.get('metadata', {})}\n"
+
+            if "conversation_history" in state_data:
+                history = state_data["conversation_history"]
+                output += f"\n💬 Conversation: {len(history)} messages\n"
+
+            return {"status": "success", "content": [{"text": output}]}
+
+        elif action == "resume":
+            # Time-travel: Load state and continue with ephemeral agent
+            if not state_file or not query:
+                return {
+                    "status": "error",
+                    "content": [{"text": "state_file and query required for resume"}],
+                }
+
+            state_path = Path(state_file).expanduser()
+            if not state_path.exists():
+                return {
+                    "status": "error",
+                    "content": [{"text": f"State file not found: {state_path}"}],
+                }
+
+            with open(state_path, "rb") as f:
+                state_data = dill.load(f)
+
+            # ✅ Create ephemeral DevDuck instance (no mutation!)
+            try:
+                from devduck import DevDuck
+
+                ephemeral_duck = DevDuck(auto_start_servers=False)
+                ephemeral_agent = ephemeral_duck.agent
+            except Exception as e:
+                return {
+                    "status": "error",
+                    "content": [{"text": f"Failed to create ephemeral DevDuck: {e}"}],
+                }
+
+            # Load saved state into ephemeral agent
+            ephemeral_agent.system_prompt = state_data["system_prompt"]
+
+            # Restore conversation history
+            if "conversation_history" in state_data:
+                saved_history = state_data["conversation_history"]
+
+                if hasattr(ephemeral_agent, "conversation_history"):
+                    ephemeral_agent.conversation_history = saved_history
+                elif hasattr(ephemeral_agent, "messages"):
+                    ephemeral_agent.messages = saved_history
+
+            # Build continuation prompt with context
+            continuation_context = f"""
+[Resumed from state: {state_path.name}]
+[Original timestamp: {state_data.get('timestamp')}]
+
+{query}
+"""
+            # Run ephemeral agent (parent agent unchanged!)
+            result = ephemeral_agent(continuation_context)
+
+            return {
+                "status": "success",
+                "content": [{"text": f"🔄 Resumed from {state_path.name}\n\n{result}"}],
+            }
+
+        elif action == "modify":
+            # Modify state metadata
+            if not state_file:
+                return {
+                    "status": "error",
+                    "content": [{"text": "state_file required for modify"}],
+                }
+
+            state_path = Path(state_file).expanduser()
+            if not state_path.exists():
+                return {
+                    "status": "error",
+                    "content": [{"text": f"State file not found: {state_path}"}],
+                }
+
+            with open(state_path, "rb") as f:
+                state_data = dill.load(f)
+
+            # Update metadata
+            if metadata:
+                state_data["metadata"].update(metadata)
+
+            # Save back
+            with open(state_path, "wb") as f:
+                dill.dump(state_data, f)
+
+            return {
+                "status": "success",
+                "content": [
+                    {
+                        "text": f"✅ Modified {state_path.name}\n📝 New metadata: {state_data['metadata']}"
+                    }
+                ],
+            }
+
+        elif action == "delete":
+            # Delete saved state
+            if not state_file:
+                return {
+                    "status": "error",
+                    "content": [{"text": "state_file required for delete"}],
+                }
+
+            state_path = Path(state_file).expanduser()
+            if not state_path.exists():
+                return {
+                    "status": "error",
+                    "content": [{"text": f"State file not found: {state_path}"}],
+                }
+
+            state_path.unlink()
+            return {
+                "status": "success",
+                "content": [{"text": f"🗑️  Deleted {state_path.name}"}],
+            }
+
+        else:
+            return {
+                "status": "error",
+                "content": [{"text": f"Unknown action: {action}"}],
+            }
+
+    except Exception as e:
+        return {"status": "error", "content": [{"text": f"Error: {e}"}]}

devduck/tools/store_in_kb.py

@@ -0,0 +1,187 @@
+"""Tool for storing data in Bedrock Knowledge Base asynchronously."""
+
+import json
+import logging
+import os
+import threading
+import time
+import uuid
+from typing import Any
+
+import boto3
+from strands import tool
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+def _store_in_kb_background(
+    content: str, title: str, kb_id: str, region_name: str
+) -> None:
+    """Background worker function that performs the actual KB storage.
+
+    This runs in a separate thread and handles all the KB operations.
+    Whole validation is done in the main thread before calling this function.
+
+    Args:
+        content: The text content to store
+        title: The title for the content
+        kb_id: The knowledge base ID
+        region_name: The AWS region to use
+    """
+    try:
+        # Generate document ID with timestamp for traceability
+        timestamp = time.strftime("%Y%m%d_%H%M%S")
+        doc_id = f"memory_{timestamp}_{str(uuid.uuid4())[:8]}"
+
+        # Package content with title for better organization
+        content_with_metadata = {
+            "title": title,
+            "action": "store",
+            "content": content,
+        }
+
+        # Initialize Bedrock agent client
+        bedrock_agent_client = boto3.client("bedrock-agent", region_name=region_name)
+
+        # Get the data source ID associated with the knowledge base
+        data_sources = bedrock_agent_client.list_data_sources(knowledgeBaseId=kb_id)
+
+        if not data_sources.get("dataSourceSummaries"):
+            logger.error(
+                f"No data sources found for knowledge base {kb_id}, region {region_name}."
+            )
+            return
+
+        # Look for a CUSTOM data source type first, as it's required for inline content ingestion
+        data_source_id = None
+        source_type = None
+
+        for ds in data_sources["dataSourceSummaries"]:
+            # Get the data source details to check its type
+            ds_detail = bedrock_agent_client.get_data_source(
+                knowledgeBaseId=kb_id, dataSourceId=ds["dataSourceId"]
+            )
+
+            # Check if this is a CUSTOM type data source
+            if ds_detail["dataSource"]["dataSourceConfiguration"]["type"] == "CUSTOM":
+                data_source_id = ds["dataSourceId"]
+                source_type = "CUSTOM"
+                logger.debug(f"Found CUSTOM data source: {data_source_id}")
+                break
+
+        # If no CUSTOM data source found, use the first available one but log a warning
+        if not data_source_id and data_sources["dataSourceSummaries"]:
+            data_source_id = data_sources["dataSourceSummaries"][0]["dataSourceId"]
+            ds_detail = bedrock_agent_client.get_data_source(
+                knowledgeBaseId=kb_id, dataSourceId=data_source_id
+            )
+            source_type = ds_detail["dataSource"]["dataSourceConfiguration"]["type"]
+            logger.debug(
+                f"No CUSTOM data source found. Using {source_type} data source: {data_source_id}"
+            )
+
+        if not data_source_id:
+            logger.error(f"No suitable data source found for knowledge base {kb_id}.")
+            return
+
+        # Prepare document for ingestion based on the data source type
+        if source_type == "CUSTOM":
+            ingest_request = {
+                "knowledgeBaseId": kb_id,
+                "dataSourceId": data_source_id,
+                "documents": [
+                    {
+                        "content": {
+                            "dataSourceType": "CUSTOM",
+                            "custom": {
+                                "customDocumentIdentifier": {"id": doc_id},
+                                "inlineContent": {
+                                    "textContent": {
+                                        "data": json.dumps(content_with_metadata)
+                                    },
+                                    "type": "TEXT",
+                                },
+                                "sourceType": "IN_LINE",
+                            },
+                        }
+                    }
+                ],
+            }
+        elif source_type == "S3":
+            # S3 source types need a different ingestion approach
+            logger.error(
+                "S3 data source type is not supported for direct ingestion with this tool."
+            )
+            return
+        else:
+            logger.error(f"Unsupported data source type: {source_type}")
+            return
+
+        # Ingest document into knowledge base
+        _ = bedrock_agent_client.ingest_knowledge_base_documents(**ingest_request)
+
+        # Log success
+        logger.info(
+            f"Successfully ingested document into knowledge base {kb_id}: {doc_id}"
+        )
+
+    except Exception as e:
+        logger.error(f"Error ingesting into knowledge base: {e!s}")
+
+
+@tool
+def store_in_kb(
+    content: str, title: str | None = None, knowledge_base_id: str | None = None
+) -> dict[str, Any]:
+    """Store content in a Bedrock Knowledge Base using real-time ingestion.
+
+    This version runs asynchronously in a background thread and returns immediately.
+
+    Args:
+        content: The text content to store in the knowledge base.
+        title: Optional title for the content. If not provided, a timestamp will be used.
+        knowledge_base_id: Optional knowledge base ID. If not provided, will use the STRANDS_KNOWLEDGE_BASE_ID env.
+
+    Returns:
+        A dictionary containing the result of the operation.
+    """
+    # All validation done in main thread before spawning background thread
+
+    # Validate content first
+    if not content or not content.strip():
+        return {"status": "error", "content": [{"text": "❌ Content cannot be empty"}]}
+
+    # Resolve and validate knowledge base ID early (addresses environment variable race condition)
+    kb_id = knowledge_base_id or os.getenv("STRANDS_KNOWLEDGE_BASE_ID")
+    if not kb_id:
+        return {
+            "status": "error",
+            "content": [
+                {
+                    "text": "❌ No knowledge base ID provided or found in environment variables STRANDS_KNOWLEDGE_BASE_ID"
+                }
+            ],
+        }
+
+    region_name = os.getenv("AWS_REGION", "us-west-2")
+
+    doc_title = title or f"Strands Memory {time.strftime('%Y%m%d_%H%M%S')}"
+
+    thread = threading.Thread(
+        target=_store_in_kb_background,
+        args=(content, doc_title, kb_id, region_name),
+        daemon=True,
+    )
+    thread.start()
+
+    # Return immediately with status
+    return {
+        "status": "success",
+        "content": [
+            {"text": "✅ Started background task to store content in knowledge base:"},
+            {"text": f"📝 Title: {doc_title}"},
+            {"text": f"🗄️ Knowledge Base ID: {kb_id}"},
+            {"text": "⏱️ Processing in background..."},
+        ],
+    }