ai-infra 0.1.3__tar.gz

ai_infra-0.1.3/PKG-INFO

@@ -0,0 +1,39 @@
+Metadata-Version: 2.3
+Name: ai-infra
+Version: 0.1.3
+Summary: Infrastructure for efficient and scalable AI applications.
+License: MIT
+Keywords: ai,langchain,langgraph,fastapi,infra,llm,mcp
+Author: Ali Khatami
+Author-email: aliikhatami94@gmail.com
+Requires-Python: >=3.11,<4.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Framework :: FastAPI
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Typing :: Typed
+Requires-Dist: fastapi (>=0.116.1,<0.117.0)
+Requires-Dist: langchain (>=0.3.27,<0.4.0)
+Requires-Dist: langchain-anthropic (>=0.3.18,<0.4.0)
+Requires-Dist: langchain-google-genai (>=2.1.9,<3.0.0)
+Requires-Dist: langchain-mcp-adapters (>=0.1.9,<0.2.0)
+Requires-Dist: langchain-openai (>=0.3.29,<0.4.0)
+Requires-Dist: langchain-xai (>=0.2.5,<0.3.0)
+Requires-Dist: langgraph (>=0.6.4,<0.7.0)
+Requires-Dist: langsmith (>=0.4.13,<0.5.0)
+Requires-Dist: mcp[cli] (>=1.13.1,<2.0.0)
+Requires-Dist: pydantic (>=2.11,<3.0)
+Requires-Dist: python-dotenv (>=1.1.1,<2.0.0)
+Project-URL: Documentation, https://github.com/your-org/ai-infra#readme
+Project-URL: Homepage, https://github.com/your-org/ai-infra
+Project-URL: Issues, https://github.com/your-org/ai-infra/issues
+Project-URL: Repository, https://github.com/your-org/ai-infra
+Description-Content-Type: text/markdown
+
+# ai-infra
+

ai_infra-0.1.3/README.md

@@ -0,0 +1 @@
+# ai-infra

ai_infra-0.1.3/pyproject.toml

@@ -0,0 +1,59 @@
+[tool.poetry]
+name = "ai-infra"
+version = "0.1.3"
+description = "Infrastructure for efficient and scalable AI applications."
+authors = ["Ali Khatami <aliikhatami94@gmail.com>"]
+license = "MIT"
+readme = "README.md"
+packages = [{ include = "ai_infra", from = "src" }]
+keywords = ["ai", "langchain", "langgraph", "fastapi", "infra", "llm", "mcp"]
+
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Framework :: FastAPI",
+    "Typing :: Typed"
+]
+
+[tool.poetry.urls]
+Homepage = "https://github.com/your-org/ai-infra"
+Repository = "https://github.com/your-org/ai-infra"
+Issues = "https://github.com/your-org/ai-infra/issues"
+Documentation = "https://github.com/your-org/ai-infra#readme"
+
+[tool.poetry.dependencies]
+python = ">=3.11,<4.0"
+
+# Core AI/Infra
+fastapi = "^0.116.1"
+pydantic = "^2.11"
+python-dotenv = "^1.1.1"
+
+# LangChain ecosystem
+langgraph = ">=0.6.4,<0.7.0"
+langchain = "^0.3.27"
+langchain-openai = "^0.3.29"
+langchain-xai = "^0.2.5"
+langchain-anthropic = "^0.3.18"
+langchain-google-genai = "^2.1.9"
+langchain-mcp-adapters = "^0.1.9"
+langsmith = "^0.4.13"
+
+# MCP integration
+mcp = {extras = ["cli"], version = "^1.13.1"}
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.0"
+pytest-asyncio = "^0.23.0"
+ruff = "^0.5.0"
+mypy = "^1.10.0"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

ai_infra-0.1.3/src/ai_infra/__init__.py

@@ -0,0 +1,20 @@
+import os
+from dotenv import load_dotenv, find_dotenv
+
+if not os.environ.get("AI_INFRA_ENV_LOADED"):
+    load_dotenv(find_dotenv(usecwd=True))
+    os.environ["AI_INFRA_ENV_LOADED"] = "1"
+
+# Re-export primary public API components
+from ai_infra.llm.core import CoreLLM
+from ai_infra.graph.core import CoreGraph
+from ai_infra.llm.providers import Providers
+from ai_infra.llm.providers.models import Models
+
+__all__ = [
+    "CoreGraph",
+    "Models",
+    "Providers",
+    "CoreMCP",
+]
+

ai_infra-0.1.3/src/ai_infra/graph/__init__.py

@@ -0,0 +1,8 @@
+from ai_infra.graph.models import Edge, ConditionalEdge
+from ai_infra.graph.core import CoreGraph
+
+__all__ = [
+    "CoreGraph",
+    "Edge",
+    "ConditionalEdge",
+]

ai_infra-0.1.3/src/ai_infra/graph/core.py

@@ -0,0 +1,156 @@
+from collections.abc import AsyncIterator, Iterator
+from typing import Any, Sequence, Union, Dict
+from langgraph.constants import START, END
+from langgraph.graph import StateGraph
+
+from ai_infra.graph.models import GraphStructure, EdgeType
+from ai_infra.graph.utils import (
+    normalize_node_definitions, normalize_initial_state,
+    build_edges, wrap_node, normalize_stream_mode,
+    make_hook, make_trace_fn, make_trace_wrapper
+)
+
+class CoreGraph:
+    def __init__(
+            self,
+            *,
+            state_type: type,
+            node_definitions: Union[Sequence, dict],
+            edges: Sequence[EdgeType],
+            checkpointer=None,
+            store=None
+    ):
+        if not (isinstance(state_type, type) and (issubclass(state_type, dict) or hasattr(state_type, '__annotations__'))):
+            raise ValueError("state_type must be a TypedDict or dict subclass")
+        self.state_type = state_type
+
+        node_definitions = normalize_node_definitions(node_definitions)
+        self.node_definitions = list(node_definitions.items())
+
+        # centralize edge building/validation + START/END
+        regular_edges, conditional_edges = build_edges(list(node_definitions.keys()), edges)
+        self.edges = regular_edges
+        self.conditional_edges = conditional_edges
+
+        self._checkpointer = checkpointer
+        self._store = store
+        self.graph = self._build_graph().compile(checkpointer=self._checkpointer, store=self._store)
+
+    def _build_graph(self, node_items=None, sync: bool=False) -> StateGraph:
+        wf = StateGraph(self.state_type)
+        node_items = node_items or self.node_definitions
+        for name, fn in node_items:
+            wf.add_node(name, wrap_node(fn, sync))
+        for start, router_fn, path_map in self.conditional_edges:
+            wf.add_conditional_edges(start, wrap_node(router_fn, sync), path_map)
+        for start, end in self.edges:
+            wf.add_edge(start, end)
+        return wf
+
+    def _prepare_run(
+            self,
+            initial_state=None,
+            *,
+            config=None,
+            on_enter=None,
+            on_exit=None,
+            trace=None,
+            sync: bool=False,
+            **kwargs
+    ):
+        initial_state = normalize_initial_state(initial_state, kwargs)
+        # fast path: no hooks → use cached compiled graph
+        if not (on_enter or on_exit or trace):
+            compiled = self._build_graph(sync=sync).compile(checkpointer=self._checkpointer, store=self._store) if sync else self.graph
+            return compiled, initial_state, config
+
+        # else, patch & recompile
+        on_enter_fn = make_hook(on_enter, sync=sync)
+        on_exit_fn = make_hook(on_exit, sync=sync)
+        trace_fn = make_trace_fn(trace, sync=sync)
+        patched_nodes = [
+            (name, make_trace_wrapper(name, wrap_node(fn, sync), on_enter_fn, on_exit_fn, trace_fn, sync))
+            for name, fn in self.node_definitions
+        ]
+        compiled = self._build_graph(node_items=patched_nodes, sync=sync).compile(checkpointer=self._checkpointer, store=self._store)
+        return compiled, initial_state, config
+
+    # ---- invoke -----------------------------------------------------------------
+    async def arun(self, initial_state=None, *, config=None, on_enter=None, on_exit=None, trace=None, **kwargs) -> Any:
+        compiled, initial_state, config = self._prepare_run(initial_state, config=config, on_enter=on_enter, on_exit=on_exit, trace=trace, sync=False, **kwargs)
+        return await compiled.ainvoke(initial_state, config=config) if config is not None else await compiled.ainvoke(initial_state)
+
+    def run(self, initial_state=None, *, config=None, on_enter=None, on_exit=None, trace=None, **kwargs) -> Any:
+        compiled, initial_state, config = self._prepare_run(initial_state, config=config, on_enter=on_enter, on_exit=on_exit, trace=trace, sync=True, **kwargs)
+        return compiled.invoke(initial_state, config=config) if config is not None else compiled.invoke(initial_state)
+
+    # ---- streaming ---------------------------------------------------------------
+    async def astream(self, initial_state=None, *, config=None, stream_mode=("updates","values")) -> AsyncIterator[tuple[str, Any]]:
+        stream_mode = normalize_stream_mode(stream_mode)
+        compiled, initial_state, config = self._prepare_run(initial_state, config=config, sync=False)
+        async for mode, chunk in compiled.astream(initial_state, config=config, stream_mode=stream_mode):
+            yield mode, chunk
+
+    def stream(self, initial_state=None, *, config=None, stream_mode=("updates","values")) -> Iterator[tuple[str, Any]]:
+        stream_mode = normalize_stream_mode(stream_mode)
+        compiled, initial_state, config = self._prepare_run(initial_state, config=config, sync=True)
+        for mode, chunk in compiled.stream(initial_state, config=config, stream_mode=stream_mode):
+            yield mode, chunk
+
+    async def astream_values(self, initial_state=None, *, config=None):
+        async for _, chunk in self.astream(initial_state, config=config, stream_mode="values"):
+            yield chunk
+
+    def stream_values(self, initial_state=None, *, config=None):
+        for _, chunk in self.stream(initial_state, config=config, stream_mode="values"):
+            yield chunk
+
+    # ---- analysis / debug --------------------------------------------------------
+    def analyze(self) -> GraphStructure:
+        nodes = [name for name, _ in self.node_definitions]
+        entry_points = [end for start, end in self.edges if start == START] or nodes[:1]
+        exit_points = [start for start, end in self.edges if end == END]
+        conditional_edges_data = [
+            {"start": start, "router_function": getattr(router_fn, '__name__', str(router_fn)), "path_options": list(path_map.keys())}
+            for start, router_fn, path_map in self.conditional_edges
+        ] if self.conditional_edges else None
+        state_schema = {key: getattr(value, '__name__', str(value)) for key, value in getattr(self.state_type, '__annotations__', {}).items()}
+        # reachability
+        reachable = set(entry_points)
+        edges_map = {start: [] for start, _ in self.edges}
+        for start, end in self.edges:
+            edges_map.setdefault(start, []).append(end)
+        queue = list(entry_points)
+        while queue:
+            node = queue.pop(0)
+            for nbr in edges_map.get(node, []):
+                if nbr not in reachable and nbr not in (START, END):
+                    reachable.add(nbr); queue.append(nbr)
+        unreachable = [n for n in nodes if n not in reachable]
+
+        return GraphStructure(
+            state_type_name=self.state_type.__name__,
+            state_schema=state_schema,
+            node_count=len(nodes),
+            nodes=nodes,
+            edge_count=len(self.edges),
+            edges=self.edges,
+            conditional_edge_count=len(self.conditional_edges) if self.conditional_edges else 0,
+            conditional_edges=conditional_edges_data,
+            entry_points=entry_points,
+            exit_points=exit_points,
+            has_memory=self._checkpointer is not None,
+            unreachable=unreachable
+        )
+
+    def describe(self) -> Dict:
+        return self.analyze().model_dump()
+
+    def get_state(self, config):
+        return self.graph.get_state(config)
+
+    def get_state_history(self, config):
+        return list(self.graph.get_state_history(config))
+
+    def get_arch_diagram(self) -> str:
+        return self.graph.get_graph().draw_mermaid()

ai_infra-0.1.3/src/ai_infra/graph/examples/01_graph_basic.py

@@ -0,0 +1,45 @@
+"""01_graph_basic: Basic graph with conditional looping.
+Usage: python -m quickstart.run graph_basic
+"""
+from typing_extensions import TypedDict
+from langgraph.graph import END
+from ai_infra.graph.core import CoreGraph
+from ai_infra.graph.models import Edge, ConditionalEdge
+
+MAX_VALUE = 40
+
+class MyState(TypedDict):
+    value: int
+
+def inc(state: MyState) -> MyState:
+    """Increment value."""
+    state["value"] += 1
+    return state
+
+def mul(state: MyState) -> MyState:
+    """Double value."""
+    state["value"] *= 2
+    return state
+
+
+def _trace(node_name, state, event):  # type: ignore[override]
+    print(f"{event.upper()} node={node_name} state={state}")
+
+
+graph = CoreGraph(
+    state_type=MyState,
+    node_definitions=[inc, mul],
+    edges=[
+        Edge(start="inc", end="mul"),
+        ConditionalEdge(
+            start="mul",
+            router_fn=lambda s: "inc" if s["value"] < MAX_VALUE else END,
+            targets=["inc", END],
+        ),
+    ],
+)
+
+
+def main():
+    result = graph.run({"value": 1}, trace=_trace)
+    print("Final:", result)

ai_infra-0.1.3/src/ai_infra/graph/examples/02_graph_stream_values.py

@@ -0,0 +1,27 @@
+"""02_graph_stream_values: Stream only state value snapshots.
+Usage: python -m quickstart.run graph_stream_values
+"""
+from typing_extensions import TypedDict
+from ai_infra.graph.core import CoreGraph
+from ai_infra.graph.models import Edge
+
+class MyState(TypedDict):
+    value: int
+
+def inc(state: MyState) -> MyState:
+    state["value"] += 1
+    return state
+
+def main():
+    graph = CoreGraph(
+        state_type=MyState,
+        node_definitions=[inc],
+        edges=[Edge(start="inc", end="inc")],  # simple loop; rely on user to break (example)
+    )
+    # For demonstration, manually break after 5 iterations
+    iterations = 0
+    for snapshot in graph.stream_values({"value": 0}):
+        print(snapshot)
+        iterations += 1
+        if iterations >= 5:
+            break

ai_infra-0.1.3/src/ai_infra/graph/models.py

@@ -0,0 +1,35 @@
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from pydantic import BaseModel, ConfigDict
+
+class GraphStructure(BaseModel):
+    """Pydantic model representing the graph's structural information."""
+    state_type_name: str
+    state_schema: Dict[str, str]
+    node_count: int
+    nodes: List[str]
+    edge_count: int
+    edges: List[Tuple[str, str]]
+    conditional_edge_count: int
+    conditional_edges: Optional[List[Dict[str, Any]]] = None
+    entry_points: List[str]
+    exit_points: List[str]
+    has_memory: bool
+    unreachable: Optional[List[str]] = None
+
+class CoreGraphConfig(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    node_definitions: Sequence[Any]
+    edges: Sequence[Tuple[str, str]]
+    conditional_edges: Optional[Sequence[Tuple[str, Any, dict]]] = None
+    memory_store: Optional[object] = None
+
+class Edge(BaseModel):
+    start: str
+    end: str
+
+class ConditionalEdge(BaseModel):
+    start: str
+    router_fn: Callable
+    targets: list[str]
+
+EdgeType = Union[Edge, ConditionalEdge]

ai_infra-0.1.3/src/ai_infra/graph/utils.py

@@ -0,0 +1,139 @@
+import inspect
+import asyncio
+from typing import Sequence, Any
+from langgraph.constants import START, END
+from ai_infra.graph.models import Edge, ConditionalEdge
+
+def normalize_node_definitions(node_definitions):
+    if isinstance(node_definitions, dict):
+        return node_definitions.copy()
+    return {fn.__name__: fn for fn in node_definitions}
+
+def normalize_initial_state(initial_state, kwargs):
+    if initial_state is None:
+        return kwargs
+    if kwargs:
+        raise ValueError("Provide either initial_state or keyword arguments, not both.")
+    return initial_state
+
+def validate_edges(edges, all_nodes):
+    for start, end in edges:
+        for endpoint in (start, end):
+            if endpoint not in all_nodes and endpoint not in (START, END):
+                raise ValueError(f"Edge endpoint '{endpoint}' is not a known node or START/END")
+
+def validate_conditional_edges(conditional_edges, all_nodes):
+    for start, router_fn, path_map in conditional_edges:
+        if start not in all_nodes and start not in (START, END):
+            raise ValueError(f"Conditional edge start '{start}' is not a known node or START/END")
+        for target in path_map.values():
+            if target not in all_nodes and target not in (START, END):
+                raise ValueError(f"Conditional path target '{target}' is not a known node or START/END")
+
+def make_router_wrapper(fn, valid_targets):
+    async def wrapper(state):
+        result = await fn(state) if inspect.iscoroutinefunction(fn) else fn(state)
+        if result not in valid_targets:
+            raise ValueError(f"Router function returned '{result}', which is not in targets {valid_targets}")
+        return result
+    return wrapper
+
+def make_hook(hook, event=None, sync=False):
+    if not hook:
+        return None
+    if inspect.iscoroutinefunction(hook):
+        if sync:
+            def sync_hook(node, state):
+                return asyncio.run(hook(node, state) if event is None else hook(node, state, event))
+            return sync_hook
+        return lambda node, state: hook(node, state) if event is None else hook(node, state, event)
+    async def async_hook(node, state):
+        return hook(node, state) if event is None else hook(node, state, event)
+    return async_hook
+
+def make_trace_fn(trace, sync=False):
+    if not trace:
+        return None
+    if sync:
+        def trace_sync(node, state, event):
+            return asyncio.run(trace(node, state, event)) if inspect.iscoroutinefunction(trace) else trace(node, state, event)
+        return trace_sync
+    async def trace_async(node, state, event):
+        if inspect.iscoroutinefunction(trace):
+            await trace(node, state, event)
+        else:
+            trace(node, state, event)
+    return trace_async
+
+def make_trace_wrapper(name, fn, on_enter, on_exit, trace, sync):
+    if sync:
+        def wrapped(state):
+            if on_enter: on_enter(name, state)
+            if trace: trace(name, state, "enter")
+            result = fn(state)
+            if on_exit: on_exit(name, result)
+            if trace: trace(name, result, "exit")
+            return result
+        return wrapped
+    async def wrapped(state):
+        if on_enter: await on_enter(name, state)
+        if trace: await trace(name, state, "enter")
+        result = await fn(state)
+        if on_exit: await on_exit(name, result)
+        if trace: await trace(name, result, "exit")
+        return result
+    return wrapped
+
+# ---- new helpers ---------------------------------------------------------------
+
+def normalize_stream_mode(stream_mode):
+    if stream_mode is None:
+        return ["updates"]
+    if isinstance(stream_mode, str):
+        return [stream_mode]
+    return list(stream_mode)
+
+def wrap_node(fn, sync: bool):
+    if sync:
+        if not inspect.iscoroutinefunction(fn):
+            return fn
+        def sync_wrapper(*args, **kwargs):
+            try:
+                loop = asyncio.get_running_loop()
+                if loop.is_running():
+                    raise RuntimeError(
+                        "CoreGraph.run/stream cannot execute async nodes inside a running event loop. "
+                        "Use arun/astream instead."
+                    )
+            except RuntimeError:
+                # no running loop; safe to asyncio.run
+                pass
+            return asyncio.run(fn(*args, **kwargs))
+        return sync_wrapper
+    if inspect.iscoroutinefunction(fn):
+        return fn
+    async def async_wrapper(*args, **kwargs):
+        return fn(*args, **kwargs)
+    return async_wrapper
+
+def build_edges(node_names: Sequence[str], edges: Sequence[Any]):
+    """Return (regular_edges, conditional_edges) normalized + auto START/END guarded."""
+    all_nodes = set(node_names)
+    regular_edges, conditional_edges = [], []
+    for edge in edges:
+        if isinstance(edge, Edge):
+            regular_edges.append((edge.start, edge.end))
+        elif isinstance(edge, ConditionalEdge):
+            for target in edge.targets:
+                if target not in all_nodes and target not in (START, END):
+                    raise ValueError(f"ConditionalEdge target '{target}' is not a known node or START/END")
+            conditional_edges.append((edge.start, make_router_wrapper(edge.router_fn, edge.targets), {t: t for t in edge.targets}))
+        else:
+            raise ValueError(f"Unknown edge type: {edge}")
+    if regular_edges and not any(s == START for s, _ in regular_edges):
+        regular_edges = [(START, regular_edges[0][0]), *regular_edges]
+    if regular_edges and not any(e == END for _, e in regular_edges):
+        regular_edges = [*regular_edges, (regular_edges[-1][1], END)]
+    validate_edges(regular_edges, all_nodes)
+    validate_conditional_edges(conditional_edges, all_nodes)
+    return regular_edges, conditional_edges

ai_infra-0.1.3/src/ai_infra/llm/__init__.py

@@ -0,0 +1,13 @@
+from ai_infra.llm.core import CoreLLM, CoreAgent, BaseLLMCore
+from ai_infra.llm.utils.settings import ModelSettings
+from ai_infra.llm.providers import Providers
+from ai_infra.llm.providers.models import Models
+
+
+__all__ = [
+    "CoreLLM",
+    "CoreAgent",
+    "ModelSettings",
+    "Models",
+    "Providers",
+]