shortgraph 0.1.0__py3-none-any.whl

shortgraph/__init__.py

@@ -0,0 +1,38 @@
+# ShortGraph Core Package
+
+# Core Graph Dimensions (LangGraph-aligned)
+from shortgraph.graph.state import StateGraph, START, END
+from shortgraph.pregel.main import Pregel
+from shortgraph.channels.manager import AgentState
+from shortgraph.checkpoint.base import BaseCheckpoint, MemoryCheckpoint
+
+# Paradigms
+from shortgraph.paradigms.smart_agent import SmartAgent
+
+# Models
+from shortgraph.models.base import BaseLLM
+from shortgraph.models.factory import LLMFactory
+
+# Tools
+from shortgraph.tools.base import Tool, tool
+from shortgraph.prebuilt.tool_node import ToolNode
+
+# Memory
+from shortgraph.memory.buffer import WindowBufferMemory
+
+__all__ = [
+    "StateGraph",
+    "Pregel",
+    "START",
+    "END",
+    "AgentState",
+    "BaseCheckpoint",
+    "MemoryCheckpoint",
+    "SmartAgent",
+    "BaseLLM",
+    "LLMFactory",
+    "Tool",
+    "tool",
+    "ToolNode",
+    "WindowBufferMemory",
+]

shortgraph/channels/__init__.py

@@ -0,0 +1,3 @@
+from .manager import AgentState
+
+__all__ = ["AgentState"]

shortgraph/channels/manager.py

@@ -0,0 +1,92 @@
+from typing import Dict, Any, List, Optional, Callable, Type
+from dataclasses import dataclass, field, asdict
+import datetime
+import operator
+import copy
+
+def default_reducer(current: Any, new: Any) -> Any:
+    """Default: Overwrite."""
+    return new
+
+def append_reducer(current: list, new: list) -> list:
+    """Append new items to current list."""
+    if current is None:
+        current = []
+    if new is None:
+        return current
+    return current + new
+
+@dataclass
+class Channel:
+    """
+    Defines how a specific field in the state should be updated.
+    """
+    name: str
+    reducer: Callable[[Any, Any], Any] = default_reducer
+    default: Any = None
+
+class AgentState:
+    """
+    Enhanced AgentState with reducer support.
+    """
+    def __init__(self, data: Optional[Dict[str, Any]] = None, channels: Optional[Dict[str, Channel]] = None):
+        self._data = data or {}
+        self._channels = channels or {}
+        
+        # Ensure default channels exist
+        if "messages" not in self._channels:
+            self._channels["messages"] = Channel("messages", reducer=append_reducer, default=[])
+        if "history" not in self._channels:
+            self._channels["history"] = Channel("history", reducer=append_reducer, default=[])
+        
+        # Initialize defaults
+        for name, channel in self._channels.items():
+            if name not in self._data:
+                self._data[name] = channel.default
+
+    def update(self, updates: Dict[str, Any]):
+        """
+        Apply updates using the defined reducers.
+        """
+        for key, value in updates.items():
+            if key in self._channels:
+                current_val = self._data.get(key, self._channels[key].default)
+                self._data[key] = self._channels[key].reducer(current_val, value)
+            else:
+                # Default behavior for unknown keys: overwrite
+                self._data[key] = value
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return self._data.get(key, default)
+
+    def __getitem__(self, key: str) -> Any:
+        return self._data[key]
+        
+    def __setitem__(self, key: str, value: Any):
+        self.update({key: value})
+
+    def to_dict(self) -> Dict[str, Any]:
+        return copy.deepcopy(self._data)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'AgentState':
+        return cls(data=data)
+
+    # Helper methods for backward compatibility
+    @property
+    def messages(self) -> List[Dict[str, Any]]:
+        return self.get("messages", [])
+
+    @property
+    def history(self) -> List[str]:
+        return self.get("history", [])
+
+    def add_message(self, role: str, content: str):
+        msg = {"role": role, "content": content, "timestamp": datetime.datetime.now().isoformat()}
+        self.update({"messages": [msg]}) # List for append_reducer
+
+    def set_variable(self, key: str, value: Any):
+        self.update({key: value})
+
+    def get_variable(self, key: str, default: Any = None) -> Any:
+        return self.get(key, default)

shortgraph/checkpoint/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseCheckpoint, MemoryCheckpoint
+
+__all__ = ["BaseCheckpoint", "MemoryCheckpoint"]

shortgraph/checkpoint/base.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional
+from dataclasses import dataclass, field
+import json
+import copy
+
+@dataclass
+class Snapshot:
+    state: Dict[str, Any]
+    next_node: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+class BaseCheckpoint(ABC):
+    """
+    Abstract base class for saving and loading graph state.
+    """
+    
+    @abstractmethod
+    def put(self, thread_id: str, checkpoint_id: str, snapshot: Dict[str, Any]) -> None:
+        """Save a snapshot (state + execution metadata)."""
+        pass
+
+    @abstractmethod
+    def get(self, thread_id: str, checkpoint_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve a snapshot. 
+        """
+        pass
+
+    @abstractmethod
+    def list_history(self, thread_id: str) -> Dict[str, Any]:
+        """List all checkpoints for a thread."""
+        pass
+
+class MemoryCheckpoint(BaseCheckpoint):
+    """
+    In-memory implementation.
+    """
+    def __init__(self):
+        # Structure: {thread_id: {checkpoint_id: snapshot_dict}}
+        self.storage: Dict[str, Dict[str, Any]] = {}
+        # Track order: {thread_id: [checkpoint_id_1, checkpoint_id_2]}
+        self.order: Dict[str, list] = {}
+
+    def put(self, thread_id: str, checkpoint_id: str, snapshot: Dict[str, Any]) -> None:
+        if thread_id not in self.storage:
+            self.storage[thread_id] = {}
+            self.order[thread_id] = []
+        
+        self.storage[thread_id][checkpoint_id] = copy.deepcopy(snapshot)
+        self.order[thread_id].append(checkpoint_id)
+
+    def get(self, thread_id: str, checkpoint_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
+        if thread_id not in self.storage:
+            return None
+        
+        if checkpoint_id is None:
+            # Get latest
+            if not self.order[thread_id]:
+                return None
+            latest_id = self.order[thread_id][-1]
+            return copy.deepcopy(self.storage[thread_id][latest_id])
+        
+        return copy.deepcopy(self.storage[thread_id].get(checkpoint_id))
+
+    def list_history(self, thread_id: str) -> list:
+        if thread_id not in self.order:
+            return []
+        return list(self.order[thread_id])
+
+class FileCheckpoint(BaseCheckpoint):
+    """
+    Simple file-based checkpointing.
+    """
+    def __init__(self, filepath: str = "checkpoints.json"):
+        self.filepath = filepath
+        self._load()
+
+    def _load(self):
+        try:
+            with open(self.filepath, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                self.storage = data.get("storage", {})
+                self.order = data.get("order", {})
+        except FileNotFoundError:
+            self.storage = {}
+            self.order = {}
+
+    def _save(self):
+        with open(self.filepath, 'w', encoding='utf-8') as f:
+            json.dump({"storage": self.storage, "order": self.order}, f, indent=2, default=str)
+
+    def put(self, thread_id: str, checkpoint_id: str, snapshot: Dict[str, Any]) -> None:
+        if thread_id not in self.storage:
+            self.storage[thread_id] = {}
+            self.order[thread_id] = []
+        
+        self.storage[thread_id][checkpoint_id] = snapshot 
+        self.order[thread_id].append(checkpoint_id)
+        self._save()
+
+    def get(self, thread_id: str, checkpoint_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
+        if thread_id not in self.storage:
+            return None
+        
+        if checkpoint_id is None:
+            if not self.order[thread_id]:
+                return None
+            latest_id = self.order[thread_id][-1]
+            return self.storage[thread_id][latest_id]
+        
+        return self.storage[thread_id].get(checkpoint_id)
+
+    def list_history(self, thread_id: str) -> list:
+        return self.order.get(thread_id, [])

shortgraph/graph/__init__.py

@@ -0,0 +1,3 @@
+from .state import StateGraph, START, END
+
+__all__ = ["StateGraph", "START", "END"]

shortgraph/graph/state.py

@@ -0,0 +1,75 @@
+from typing import Dict, Any, List, Callable, Union, Optional
+import time
+from shortgraph.channels.manager import AgentState
+from shortgraph.checkpoint.base import BaseCheckpoint
+from shortgraph.pregel.main import Pregel
+
+# Special constants
+START = "__start__"
+END = "__end__"
+
+class StateGraph:
+    """
+    A lightweight graph engine for defining agent workflows.
+    Structure definition only.
+    """
+    def __init__(self):
+        self.nodes: Dict[str, Callable] = {}
+        self.edges: Dict[str, str] = {}
+        self.conditional_edges: Dict[str, Callable] = {}
+        self.conditional_edge_metadata: Dict[str, Dict[str, str]] = {} # Store path_map for visualization
+        self.entry_point: str = ""
+        self.node_metadata: Dict[str, Dict[str, Any]] = {}
+
+    def add_node(self, name: str, func: Union[Callable, 'Pregel'], retry: int = 0):
+        """
+        Add a node. Can be a function or another Pregel graph (SubGraph).
+        :param retry: Number of times to retry if the node raises an exception.
+        """
+        self.nodes[name] = func
+        self.node_metadata[name] = {"retry": retry}
+
+    def add_edge(self, start_key: str, end_key: str):
+        if start_key == END:
+            raise ValueError("END node cannot have outgoing edges")
+        self.edges[start_key] = end_key
+
+    def add_conditional_edges(
+        self,
+        source_key: str,
+        condition_func: Callable[[AgentState], str],
+        path_map: Optional[Dict[str, str]] = None
+    ):
+        if source_key == END:
+            raise ValueError("END node cannot have outgoing edges")
+            
+        def router(state: AgentState) -> str:
+            result = condition_func(state)
+            if path_map:
+                return path_map.get(result, result)
+            return result
+            
+        self.conditional_edges[source_key] = router
+        if path_map:
+            self.conditional_edge_metadata[source_key] = path_map
+
+    def set_entry_point(self, key: str):
+        self.entry_point = key
+
+    def compile(
+        self, 
+        checkpointer: Optional[BaseCheckpoint] = None,
+        interrupt_before: Optional[List[str]] = None,
+        interrupt_after: Optional[List[str]] = None
+    ) -> Pregel:
+        if not self.entry_point:
+            raise ValueError("Graph must have an entry point.")
+        
+        pregel_app = Pregel(
+            self, 
+            checkpointer, 
+            interrupt_before=interrupt_before, 
+            interrupt_after=interrupt_after
+        )
+        
+        return pregel_app

shortgraph/memory/__init__.py

@@ -0,0 +1,3 @@
+from .buffer import WindowBufferMemory
+
+__all__ = ["WindowBufferMemory"]

shortgraph/memory/buffer.py

@@ -0,0 +1,60 @@
+from typing import List, Dict, Any, Optional
+from abc import ABC, abstractmethod
+
+class BaseMemory(ABC):
+    """
+    Abstract base class for memory management.
+    """
+    @abstractmethod
+    def load(self, state: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Load messages from state and prepare them for LLM (e.g., trimming)."""
+        pass
+
+    @abstractmethod
+    def save(self, state: Dict[str, Any], new_messages: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Save new messages to state."""
+        pass
+
+class WindowBufferMemory(BaseMemory):
+    """
+    Keeps a sliding window of the most recent k messages.
+    Preserves the System message if present.
+    """
+    def __init__(self, window_size: int = 10):
+        self.window_size = window_size
+
+    def load(self, state_messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Returns the most recent k messages.
+        Always keeps the first message if it's a 'system' message.
+        """
+        if not state_messages:
+            return []
+            
+        if len(state_messages) <= self.window_size:
+            return state_messages
+            
+        # Check for system message
+        result = []
+        start_idx = 0
+        if state_messages[0].get("role") == "system":
+            result.append(state_messages[0])
+            start_idx = 1
+            
+        # Calculate how many to take from the end
+        remaining_slots = self.window_size - len(result)
+        if remaining_slots > 0:
+            # Take the last 'remaining_slots' messages
+            recent = state_messages[-remaining_slots:]
+            # Ensure we don't duplicate if overlap (though logic above handles it)
+            result.extend(recent)
+            
+        return result
+
+    def save(self, current_messages: List[Dict[str, Any]], new_messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        In ShortGraph, state updates are usually handled by Reducers (append).
+        So this method might be used if we were manually managing state, 
+        but usually we just return what needs to be added.
+        """
+        return new_messages

shortgraph/models/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseLLM
+from .factory import LLMFactory
+from .openai import OpenAILLM
+
+__all__ = ["BaseLLM", "LLMFactory", "OpenAILLM"]

shortgraph/models/base.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional, Union
+
+class BaseLLM(ABC):
+    """
+    Abstract base class for LLMs to ensure neutrality.
+    """
+    
+    @abstractmethod
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        """
+        Generate text based on prompt.
+        """
+        pass
+
+    @abstractmethod
+    def chat(self, messages: List[Dict[str, str]], stop: Optional[List[str]] = None) -> str:
+        """
+        Chat completion interface.
+        messages format: [{"role": "user", "content": "hello"}, ...]
+        """
+        pass
+
+class MockLLM(BaseLLM):
+    """
+    A mock LLM for testing purposes.
+    """
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        return f"Mock response to: {prompt[:20]}..."
+
+    def chat(self, messages: List[Dict[str, str]], stop: Optional[List[str]] = None) -> str:
+        last_msg = messages[-1]['content']
+        return f"Mock chat response to: {last_msg[:20]}..."

shortgraph/models/factory.py

@@ -0,0 +1,102 @@
+from typing import Optional, Any
+import os
+from shortgraph.models.base import BaseLLM
+from shortgraph.models.openai import OpenAILLM
+
+class LLMFactory:
+    """
+    A simple factory to create LLM instances for various providers.
+    """
+
+    @staticmethod
+    def create_openai(
+        api_key: Optional[str] = None, 
+        model: str = "gpt-3.5-turbo"
+    ) -> BaseLLM:
+        """Create a standard OpenAI LLM instance."""
+        return OpenAILLM(api_key=api_key, model=model)
+
+    @staticmethod
+    def create_deepseek(
+        api_key: Optional[str] = None, 
+        model: str = "deepseek-chat"
+    ) -> BaseLLM:
+        """
+        Create a DeepSeek LLM instance.
+        Default Base URL: https://api.deepseek.com
+        """
+        return OpenAILLM(
+            api_key=api_key or os.environ.get("DEEPSEEK_API_KEY"),
+            base_url="https://api.deepseek.com",
+            model=model
+        )
+
+    @staticmethod
+    def create_moonshot(
+        api_key: Optional[str] = None, 
+        model: str = "moonshot-v1-8k"
+    ) -> BaseLLM:
+        """
+        Create a Moonshot (Kimi) LLM instance.
+        Default Base URL: https://api.moonshot.cn/v1
+        """
+        return OpenAILLM(
+            api_key=api_key or os.environ.get("MOONSHOT_API_KEY"),
+            base_url="https://api.moonshot.cn/v1",
+            model=model
+        )
+
+    @staticmethod
+    def create_zhipu(
+        api_key: Optional[str] = None,
+        model: str = "glm-4"
+    ) -> BaseLLM:
+        """
+        Create a ZhipuAI (ChatGLM) LLM instance.
+        Default Base URL: https://open.bigmodel.cn/api/paas/v4/
+        """
+        return OpenAILLM(
+            api_key=api_key or os.environ.get("ZHIPU_API_KEY"),
+            base_url="https://open.bigmodel.cn/api/paas/v4/",
+            model=model
+        )
+
+    @staticmethod
+    def create_ollama(
+        base_url: str = "http://localhost:11434/v1",
+        model: str = "llama3"
+    ) -> BaseLLM:
+        """
+        Create a local Ollama LLM instance (via OpenAI compatibility).
+        """
+        return OpenAILLM(
+            api_key="ollama", # API key is required but ignored by Ollama
+            base_url=base_url,
+            model=model
+        )
+
+    @staticmethod
+    def create_custom(
+        api_key: str,
+        base_url: str,
+        model: str
+    ) -> BaseLLM:
+        """Create a custom OpenAI-compatible LLM instance."""
+        return OpenAILLM(api_key=api_key, base_url=base_url, model=model)
+
+    @staticmethod
+    def create_mock() -> BaseLLM:
+        """Create a Mock LLM for testing/demo purposes."""
+        class MockSmartLLM(BaseLLM):
+            def generate(self, prompt: str, stop: Optional[list] = None) -> str:
+                # Simple logic to simulate weather tool usage
+                if "Observation:" in prompt:
+                    return "Final Answer: 根据天气预报, 北京今天天气不错, 晴转多云, 适合出行!"
+                return 'Thought: 用户想查询北京的天气, 我需要调用 get_weather 工具.\nAction: get_weather\nAction Input: "北京"'
+            
+            def chat(self, messages: list, stop: Optional[list] = None) -> str:
+                # Fallback to generate for chat interface
+                prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
+                return self.generate(prompt)
+        
+        return MockSmartLLM()

shortgraph/models/openai.py

@@ -0,0 +1,49 @@
+from typing import List, Dict, Optional, Any
+import os
+try:
+    from openai import OpenAI
+except ImportError:
+    OpenAI = None
+
+from shortgraph.models.base import BaseLLM
+
+class OpenAILLM(BaseLLM):
+    """
+    Adapter for OpenAI's GPT models and OpenAI-compatible APIs (DeepSeek, Moonshot, etc.).
+    """
+    def __init__(
+        self, 
+        api_key: Optional[str] = None, 
+        model: str = "gpt-3.5-turbo",
+        base_url: Optional[str] = None
+    ):
+        if OpenAI is None:
+            raise ImportError("Please install openai package: `pip install openai`")
+        
+        self.client = OpenAI(
+            api_key=api_key or os.environ.get("OPENAI_API_KEY"),
+            base_url=base_url or os.environ.get("OPENAI_BASE_URL")
+        )
+        self.model = model
+
+    def generate(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                stop=stop
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return f"Error generating response: {str(e)}"
+
+    def chat(self, messages: List[Dict[str, str]], stop: Optional[List[str]] = None) -> str:
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                stop=stop
+            )
+            return response.choices[0].message.content
+        except Exception as e:
+            return f"Error chatting: {str(e)}"

shortgraph/paradigms/__init__.py

@@ -0,0 +1,6 @@
+from .smart_agent import SmartAgent
+from .react import ReActAgent
+from .plan_execute import PlanExecuteAgent
+from .graph_react import GraphReActAgent
+
+__all__ = ["SmartAgent", "ReActAgent", "PlanExecuteAgent", "GraphReActAgent"]