veri-agents-api 0.1.1__py3-none-any.whl

veri_agents_api/fastapi/thread/__init__.py

@@ -0,0 +1,2 @@
+from .router import *
+from .schema import *

veri_agents_api/fastapi/thread/router.py

@@ -0,0 +1,334 @@
+import asyncio
+import json
+import logging
+from typing import Any, AsyncGenerator, Dict, List, Tuple, Callable
+from uuid import uuid4
+
+from fastapi import HTTPException, Request, APIRouter
+from fastapi.responses import StreamingResponse
+from langchain_core.callbacks import AsyncCallbackHandler
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph.graph import CompiledGraph
+
+from .schema import (
+    ChatMessage,
+    StreamInput,
+    InvokeInput,
+)
+from veri_agents_api.threads_util import ThreadInfo, ThreadsCheckpointerUtil
+from veri_agents_api.util.awaitable import as_awaitable, MaybeAwaitable
+
+log = logging.getLogger(__name__)
+
+class TokenQueueStreamingHandler(AsyncCallbackHandler):
+    """LangChain callback handler for streaming LLM tokens to an asyncio queue."""
+
+    def __init__(self, queue: asyncio.Queue):
+        self.queue = queue
+
+    async def on_llm_new_token(self, token: str, **kwargs) -> None:
+        if token:
+            await self.queue.put(token)
+
+def create_thread_router(
+    get_graph: Callable[[Request], MaybeAwaitable[CompiledGraph]],
+    get_thread_id: Callable[[Request], MaybeAwaitable[str]],
+    allow_access_thread: Callable[[str, ThreadInfo | None, Request], MaybeAwaitable[bool]] = lambda thread_id, thread_info, request: True,
+    allow_invoke_thread: Callable[[str, ThreadInfo | None, InvokeInput, Request], MaybeAwaitable[bool]] = lambda thread_id, thread_info, invoke_input, request: True,
+    invoke_runnable_config: Callable[[str, ThreadInfo | None, InvokeInput, Request], MaybeAwaitable[RunnableConfig | None]] = lambda thread_id, thread_info, invoke_input, request: None,
+    on_invoke_thread: Callable[[str, ThreadInfo | None, InvokeInput, Request], MaybeAwaitable[None]] = lambda thread_id, thread_info, invoke_input, request: None,
+    # InvokeInputCls: Type[InvokeInput] = InvokeInput,
+    **router_kwargs
+):
+    """
+    POST /invoke
+    POST /stream
+    GET /history
+    GET /feedback
+    POST /feedback
+    """
+
+    router = APIRouter(**router_kwargs)
+
+    def _parse_input(user_input: InvokeInput, thread_id: str, invoke_recvd_runnable_config: RunnableConfig | None) -> Tuple[Dict[str, Any], str]:
+        run_id = uuid4()
+        input_message = ChatMessage(type="human", content=user_input.message)
+
+        runnable_config = invoke_recvd_runnable_config or RunnableConfig()
+
+        runnable_config["configurable"] = {
+            **{
+                # used by checkpointer
+                "thread_id": thread_id,
+
+                "_has_threadinfo": True,
+
+                # "args": user_input.args,
+            }, 
+            **(runnable_config.get("configurable", {}))
+        }
+
+        kwargs = dict(
+            input={"messages": [input_message.to_langchain()]},
+            config=runnable_config
+        )
+        return kwargs, str(run_id)
+
+    @router.post("/invoke")
+    async def invoke(invoke_input: InvokeInput, request: Request) -> ChatMessage:
+        """
+        Invoke the agent with user input to retrieve a final response.
+
+        Use thread_id to persist and continue a multi-turn conversation. run_id kwarg
+        is also attached to messages for recording feedback.
+        """
+
+        graph = await as_awaitable(get_graph(request))
+        thread_id = await as_awaitable(get_thread_id(request))
+        
+
+        thread_info = await ThreadsCheckpointerUtil.get_thread_info(thread_id, graph.checkpointer)
+
+        if not await as_awaitable(allow_invoke_thread(thread_id, thread_info, invoke_input, request)):
+            raise HTTPException(status_code=403, detail="Forbidden")
+        
+        invoke_recvd_runnable_config = await as_awaitable(invoke_runnable_config(thread_id, thread_info, invoke_input, request))
+
+        kwargs, run_id = _parse_input(invoke_input, thread_id, invoke_recvd_runnable_config)
+
+        # # store this thread in the database if a new one
+        # if user_input.thread_id not in router.state.threads:
+        #     thread_info = ThreadInfo(
+        #         thread_id=user_input.thread_id,
+        #         user=principal,
+        #         workflow_id=user_input.workflow,
+        #         name=user_input.message[:50],
+        #         metadata={"router": user_input.router},
+        #     )
+        #     router.state.threads[user_input.thread_id] = thread_info
+        #     await graph.checkpointer.aput_thread(thread_info)
+
+        await as_awaitable(on_invoke_thread(thread_id, thread_info, invoke_input, request))
+
+        # langfuse_handler = CallbackHandler(
+        #     public_key=router.state.cfg.logging.langfuse.public_key,
+        #     secret_key=router.state.cfg.logging.langfuse.secret_key,
+        #     host=router.state.cfg.logging.langfuse.host,
+        #     # user_id=principal,
+        #     session_id=user_input.thread_id,
+        #     trace_name=user_input.message[:50],
+        # )
+        kwargs["config"]["callbacks"] = [] # was [langfuse_handler]
+        # kwargs["config"]["configurable"]["workflow_id"] = user_input.workflow
+        try:
+            response = await graph.ainvoke(**kwargs)
+            output = ChatMessage.from_langchain(response["messages"][-1])
+            output.run_id = str(run_id)
+            return output
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+    @router.post("/stream")
+    async def stream_agent(stream_input: StreamInput, request: Request):
+        """
+        Stream the agent's response to a user input, including intermediate messages and tokens.
+
+        Use thread_id to persist and continue a multi-turn conversation. run_id kwarg
+        is also attached to all messages for recording feedback.
+        """
+
+        graph = await as_awaitable(get_graph(request))
+        thread_id = await as_awaitable(get_thread_id(request))
+
+        thread_info = await ThreadsCheckpointerUtil.get_thread_info(thread_id, graph.checkpointer)
+
+        if not await as_awaitable(allow_invoke_thread(thread_id, thread_info, stream_input, request)):
+            raise HTTPException(status_code=403, detail="Forbidden")
+        
+        invoke_recvd_runnable_config = await as_awaitable(invoke_runnable_config(thread_id, thread_info, stream_input, request))
+
+        async def message_generator() -> AsyncGenerator[str, None]:
+            """
+            Generate a stream of messages from the agent.
+
+            This is the workhorse method for the /stream endpoint.
+            """
+            kwargs, run_id = _parse_input(stream_input, thread_id, invoke_recvd_runnable_config)
+
+            await as_awaitable(on_invoke_thread(thread_id, thread_info, stream_input, request))
+
+            # # store this thread in the database if a new one
+            # if user_input.thread_id not in router.state.threads:
+            #     thread_info = ThreadInfo(
+            #         thread_id=user_input.thread_id,
+            #         user=principal,
+            #         workflow_id=user_input.workflow,
+            #         name=user_input.message[:50],
+            #         metadata={"router": user_input.router},
+            #     )
+            #     router.state.threads[user_input.thread_id] = thread_info
+            #     await graph.checkpointer.aput_thread(thread_info)
+
+            # Use an asyncio queue to process both messages and tokens in
+            # chronological order, so we can easily yield them to the client.
+            output_queue = asyncio.Queue(maxsize=10)
+
+            # langfuse_handler = CallbackHandler(
+            #     public_key=router.state.cfg.logging.langfuse.public_key,
+            #     secret_key=router.state.cfg.logging.langfuse.secret_key,
+            #     host=router.state.cfg.logging.langfuse.host,
+            #     user_id=principal,
+            #     session_id=user_input.thread_id,
+            #     trace_name=user_input.message[:50],
+            # )
+            if stream_input.stream_tokens:
+                kwargs["config"]["callbacks"] = [
+                    TokenQueueStreamingHandler(queue=output_queue),
+                    # langfuse_handler,
+                ]
+            # kwargs["config"]["configurable"]["workflow_id"] = stream_input.workflow
+
+            # Pass the agent's stream of messages to the queue in a separate task, so
+            # we can yield the messages to the client in the main thread.
+            async def run_agent_stream():
+                async for s in graph.astream(**kwargs, stream_mode="updates"):
+                    await output_queue.put(s)
+                await output_queue.put(None)
+
+            stream_task = asyncio.create_task(run_agent_stream())
+
+            # Process the queue and yield messages over the SSE stream.
+            while s := await output_queue.get():
+                log.info("Got from queue: %s: %s", type(s), s)
+                if isinstance(s, str):
+                    # str is an LLM token
+                    yield f"data: {json.dumps({'type': 'token', 'content': s})}\n\n"
+                    continue
+
+                # Otherwise, s should be a dict of state updates for each node in the graph.
+                # s could have updates for multiple nodes, so check each for messages.
+                new_messages = []
+                for _, state in s.items():
+                    new_messages.extend(state["messages"])
+                for message in new_messages:
+                    try:
+                        chat_message = ChatMessage.from_langchain(message)
+                        chat_message.run_id = str(run_id)
+                    except Exception as e:
+                        yield f"data: {json.dumps({'type': 'error', 'content': f'Error parsing message: {e}'})}\n\n"
+                        continue
+                    # LangGraph re-sends the input message, which feels weird, so drop it
+                    if (
+                        chat_message.type == "human"
+                        and chat_message.content == stream_input.message
+                    ):
+                        continue
+                    yield f"data: {json.dumps({'type': 'message', 'content': chat_message.dict()})}\n\n"
+
+            await stream_task
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(
+            message_generator(),
+            media_type="text/event-stream",
+        )
+
+    @router.get("/history")
+    async def get_history(request: Request) -> List[ChatMessage]:
+        """
+        Get the history of a thread.
+        """
+
+        graph = await as_awaitable(get_graph(request))
+        thread_id = await as_awaitable(get_thread_id(request))
+
+        thread_info = await ThreadsCheckpointerUtil.get_thread_info(thread_id, graph.checkpointer)
+
+        if not await as_awaitable(allow_access_thread(thread_id, thread_info, request)):
+            raise HTTPException(status_code=403, detail="Forbidden")
+
+        # agent: CompiledGraph = router.state.workflows[workflow].get_graph()
+        config = RunnableConfig(configurable={
+            # used by checkpointer
+            "thread_id": thread_id,
+        })
+        state = await graph.aget_state(config)
+        messages = state.values.get("messages", [])
+
+        converted_messages: List[ChatMessage] = []
+        for message in messages:
+            try:
+                chat_message = ChatMessage.from_langchain(message)
+                converted_messages.append(chat_message)
+            except Exception as e:
+                log.error(f"Error parsing message: {e}")
+                continue
+        return converted_messages
+
+    # @router.get("/feedback")
+    # async def get_feedback(request: Request, thread_id: str):
+    #     """Get all feedback for a thread.
+    #
+    #     Arguments:
+    #         thread_id: The ID of the thread to get feedback for.
+    #     """
+    #     # if thread_id not in router.state.threads:
+    #     #     raise HTTPException(status_code=404, detail=f"Unknown thread: {thread_id}")
+    #     # assert_viewer_can_assume_identity(
+    #     #     request, principal=router.state.threads[thread_id].user
+    #     # )
+    #     feedback = [
+    #         f.model_dump(mode="json")
+    #         async for f in graph.checkpointer.alist_feedback(thread_id=thread_id)
+    #     ]
+    #     return feedback
+    #
+    # @router.post("/feedback")
+    # async def feedback(feedback: Feedback, request: Request):
+    #     """
+    #     Record feedback for a run of the agent.
+    #
+    #     Arguments:
+    #         feedback: The feedback to record.
+    #     """
+    #     if feedback.thread_id not in router.state.threads:
+    #         raise HTTPException(
+    #             status_code=404, detail=f"Unknown thread: {feedback.thread_id}"
+    #         )
+    #     assert_viewer_can_assume_identity(
+    #         request, principal=router.state.threads[feedback.thread_id].user
+    #     )
+    #
+    #     # store in database
+    #     try:
+    #         await graph.checkpointer.aput_feedback(feedback)
+    #         db_status = "success"
+    #     except Exception as e:
+    #         log.error(f"Error storing feedback in database: {e}")
+    #         db_status = "error"
+    #
+    #     ## Also store in Langfuse
+    #     ## We don't have the run_id, but need it for Langfuse
+    #     ## The run_id is currently not store in the database.
+    #     # try:
+    #     #     langfuse = Langfuse(
+    #     #         public_key=router.state.cfg.logging.langfuse.public_key,
+    #     #         secret_key=router.state.cfg.logging.langfuse.secret_key,
+    #     #         host=router.state.cfg.logging.langfuse.host,
+    #     #     )
+    #     #     langfuse.score(
+    #     #         trace_id=feedback.run_id,
+    #     #         name=feedback.key,
+    #     #         value=feedback.score,
+    #     #         comment=feedback.kwargs.get("comment", ""),
+    #     #     )
+    #     #     langfuse_status = "success"
+    #     # except Exception as e:
+    #     #     log.error(f"Error storing feedback in Langfuse: {e}")
+    #     #     langfuse_status = "error"
+    #
+    #     langfuse_status = "not implemented"
+    #
+    #     return {"db_status": db_status, "langfuse_status": langfuse_status}
+
+    return router

veri_agents_api/fastapi/thread/schema.py

@@ -0,0 +1,169 @@
+from datetime import datetime
+from typing import Dict, Any, List, Literal, Optional, Union
+from langchain_core.messages import (
+    BaseMessage,
+    HumanMessage,
+    AIMessage,
+    ToolMessage,
+    ToolCall,
+    message_to_dict,
+    messages_from_dict,
+)
+from pydantic import BaseModel, Field
+
+class InvokeInput(BaseModel):
+    """Basic user input for the agent."""
+
+    message: str = Field(
+        description="User input to the agent.",
+        examples=["What is the weather in Tokyo?"],
+    )
+    # args: Dict[str, Any] = Field(
+    #     description="Arguments to pass to the workflow.",
+    #     default={},
+    #     examples=[{"kb": "veritone_support"}],
+    # )
+    # user: Optional[str] = Field(
+    #     description="A user identifier to validate the user in knowledge bases and other tools.",
+    #     default=None,
+    #     examples=["jjohnson", "ccarlson"],
+    # )
+    # thread_id: str = Field(
+    #     description="Thread ID to persist and continue a multi-turn conversation.",
+    #     examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    # )
+
+
+class StreamInput(InvokeInput):
+    """User input for streaming the agent's response."""
+
+    stream_tokens: bool = Field(
+        description="Whether to stream LLM tokens to the client.",
+        default=True,
+    )
+
+
+class AgentResponse(BaseModel):
+    """Response from the agent when called via /invoke."""
+
+    message: Dict[str, Any] = Field(
+        description="Final response from the agent, as a serialized LangChain message.",
+        examples=[
+            {
+                "message": {
+                    "type": "ai",
+                    "data": {
+                        "content": "The weather in Tokyo is 70 degrees.",
+                        "type": "ai",
+                    },
+                }
+            }
+        ],
+    )
+
+
+class ChatMessage(BaseModel):
+    """Message in a chat."""
+
+    type: Literal["human", "ai", "tool"] = Field(
+        description="Role of the message.",
+        examples=["human", "ai", "tool"],
+    )
+    content: Union[str, list[Union[str, dict]]] = Field(
+        description="Content of the message.",
+        examples=["Hello, world!"],
+    )
+    tool_calls: List[ToolCall] = Field(
+        description="Tool calls in the message.",
+        default=[],
+    )
+    tool_call_id: str | None = Field(
+        description="Tool call that this message is responding to.",
+        default=None,
+        examples=["call_Jja7J89XsjrOLA5r!MEOW!SL"],
+    )
+    run_id: str | None = Field(
+        description="Run ID of the message.",
+        default=None,
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    original: Dict[str, Any] = Field(
+        description="Original LangChain message in serialized form.",
+        default={},
+    )
+
+    @classmethod
+    def from_langchain(cls, message: BaseMessage) -> "ChatMessage":
+        """Create a ChatMessage from a LangChain message."""
+        original = message_to_dict(message)
+        match message:
+            case HumanMessage():
+                human_message = cls(
+                    type="human", content=message.content, original=original
+                )
+                return human_message
+            case AIMessage():
+                ai_message = cls(type="ai", content=message.content, original=original)
+                if message.tool_calls:
+                    ai_message.tool_calls = message.tool_calls
+                return ai_message
+            case ToolMessage():
+                tool_message = cls(
+                    type="tool",
+                    content=message.content,
+                    tool_call_id=message.tool_call_id,
+                    original=original,
+                )
+                return tool_message
+            case _:
+                raise ValueError(
+                    f"Unsupported message type: {message.__class__.__name__}"
+                )
+
+    def to_langchain(self) -> BaseMessage:
+        """Convert the ChatMessage to a LangChain message."""
+        if self.original:
+            return messages_from_dict([self.original])[0]
+        match self.type:
+            case "human":
+                return HumanMessage(content=self.content)
+            case _:
+                raise NotImplementedError(f"Unsupported message type: {self.type}")
+
+    def pretty_print(self) -> None:
+        """Pretty print the ChatMessage."""
+        lc_msg = self.to_langchain()
+        lc_msg.pretty_print()
+
+    def get_artifact(self) -> Optional[Dict[str, Any]]:
+        """Get the artifact from the message if there is one."""
+        if (
+            self.original.get("data")
+            and self.original["data"].get("artifact")
+            and self.original["data"]["artifact"]
+        ):
+            return self.original["data"]["artifact"]
+        return None
+
+
+class Feedback(BaseModel):
+    """Feedback for a run."""
+
+    message_id: str = Field(
+        description="Message ID to record feedback for.",
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    # thread_id: str = Field(
+    #     description="Thread ID to record feedback for.",
+    #     examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    # )
+    score: float = Field(
+        description="Feedback score.",
+        examples=[0.8],
+    )
+    kwargs: Dict[str, Any] = Field(
+        description="Additional feedback kwargs, passed to LangSmith.",
+        default={},
+        examples=[{"comment": "In-line human feedback"}],
+    )
+    creation: datetime = Field(default_factory=datetime.now)

veri_agents_api/fastapi/thread/test/multi_thread.py

@@ -0,0 +1,59 @@
+from typing import Annotated
+
+from langchain_aws import ChatBedrock
+from typing_extensions import TypedDict
+
+from fastapi import FastAPI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+import uvicorn
+
+from veri_agents_api.fastapi.thread import create_thread_router
+
+if __name__ == "__main__":
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    graph_builder = StateGraph(State)
+
+    llm = ChatBedrock(model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0") # pyright: ignore[reportCallIssue]
+
+    def chatbot(state: State):
+        return {"messages": [llm.invoke(state["messages"])]}
+
+    # The first argument is the unique node name
+    # The second argument is the function or object that will be called whenever
+    # the node is used.
+    graph_builder.add_node("chatbot", chatbot)
+    graph_builder.set_entry_point("chatbot")
+    graph_builder.set_finish_point("chatbot")
+
+
+
+    # in-memory persistence
+    memory = MemorySaver()
+    graph = graph_builder.compile(checkpointer=memory)
+
+    # veri-agents convenience router
+    thread_router = create_thread_router(
+        # same graph for every request
+        get_graph=lambda req: graph,
+        # derive thread id from /thread/{thread_id} path param
+        get_thread_id=lambda req: req.path_params["thread_id"]
+    )
+
+    # root fastapi app
+    app = FastAPI()
+    app.include_router(thread_router, prefix="/threads/{thread_id}")
+
+    uvicorn.run(app, port=5000, log_level="info")
+    # you can now access:
+    #   GET  /openapi.json
+    #   GET  /threads
+    #   GET /thread/{thread_id}/info
+    #   POST /threads/{thread_id}/invoke
+    #   POST /threads/{thread_id}/stream
+    #   GET  /threads/{thread_id}/history
+    #   GET  /threads/{thread_id}/feedback
+    #   POST /threads/{thread_id}/feedback

veri_agents_api/fastapi/thread/test/single_thread.py

@@ -0,0 +1,58 @@
+from typing import Annotated
+
+from langchain_aws import ChatBedrock
+from typing_extensions import TypedDict
+
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+from fastapi import FastAPI
+import uvicorn
+
+from veri_agents_api.fastapi.thread import create_thread_router
+
+if __name__ == "__main__":
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    graph_builder = StateGraph(State)
+
+    llm = ChatBedrock(model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0") # pyright: ignore[reportCallIssue]
+
+    def chatbot(state: State):
+        return {"messages": [llm.invoke(state["messages"])]}
+
+    # The first argument is the unique node name
+    # The second argument is the function or object that will be called whenever
+    # the node is used.
+    graph_builder.add_node("chatbot", chatbot)
+    graph_builder.set_entry_point("chatbot")
+    graph_builder.set_finish_point("chatbot")
+
+
+
+
+    # in-memory persistence
+    memory = MemorySaver()
+    graph = graph_builder.compile(checkpointer=memory)
+
+    # veri-agents convenience router
+    thread_router = create_thread_router(
+        # same graph for every request
+        get_graph=lambda req: graph,
+        # same thread for every request
+        get_thread_id=lambda req: "inmem"
+    )
+
+    # root fastapi app
+    app = FastAPI()
+    app.include_router(thread_router)
+
+    uvicorn.run(app, port=5000, log_level="info")
+    # you can now access:
+    #   GET  /openapi.json
+    #   POST /invoke
+    #   POST /stream
+    #   GET  /history
+    #   GET  /feedback
+    #   POST /feedback

veri_agents_api/fastapi/threads/__init__.py

@@ -0,0 +1 @@
+from .router import *

veri_agents_api/fastapi/threads/router.py

@@ -0,0 +1,75 @@
+import logging
+from typing import Callable, cast, Type, Awaitable
+
+from fastapi import HTTPException, Request, APIRouter
+from langgraph.graph.graph import CompiledGraph
+
+from veri_agents_api.fastapi.thread import (
+    create_thread_router as create_thread_router, InvokeInput
+)
+from veri_agents_api.threads_util import ThreadsCheckpointerUtil, ThreadInfo
+from veri_agents_api.util.awaitable import MaybeAwaitable, as_awaitable
+
+log = logging.getLogger(__name__)
+
+
+def create_threads_router(
+    get_graph: Callable[[Request], MaybeAwaitable[CompiledGraph]],
+    allow_access_thread: Callable[[str, ThreadInfo | None, Request], MaybeAwaitable[bool]] = lambda thread_id, thread_info, request: True,
+    allow_invoke_thread: Callable[[str, ThreadInfo | None, InvokeInput, Request], MaybeAwaitable[bool]] = lambda thread_id, thread_info, invoke_input, request: True,
+    on_invoke_thread: Callable[[str, ThreadInfo | None, InvokeInput, Request], MaybeAwaitable[None]] = lambda thread_id, thread_info, invoke_input, request: None,
+    # InvokeInputCls: Type[InvokeInput] = InvokeInput,
+    **router_kwargs
+):
+    router = APIRouter(prefix="/threads", **router_kwargs)
+
+    thread_router = create_thread_router(
+        # derive thread id from /thread/{thread_id} path param
+        get_thread_id=lambda req: req.path_params["thread_id"],
+
+        # arg passthrough - TODO: make more elegant
+        get_graph=get_graph,
+        allow_access_thread=allow_access_thread,
+        allow_invoke_thread=allow_invoke_thread,
+        on_invoke_thread=on_invoke_thread,
+
+        # InvokeInputCls=InvokeInputCls
+    )
+
+    @router.get("/")
+    async def get_threads(request: Request):
+        """Get all threads the user has access to."""
+
+        graph = await as_awaitable(get_graph(request))
+
+        all_thread_info = await ThreadsCheckpointerUtil.list_threads(graph.checkpointer)
+
+        accessible_thread_info: list[ThreadInfo] = []
+        for thread_info in all_thread_info:
+            if allow_access_thread(thread_info.thread_id, thread_info, request):
+                accessible_thread_info.append(thread_info)
+
+        return accessible_thread_info
+
+    @router.get("/{thread_id}/info")
+    async def get_thread_by_id(thread_id: str, request: Request):
+        """Get a thread by its ID.
+
+        Arguments:
+            thread_id: The ID of the thread to get.
+        """
+        graph = await as_awaitable(get_graph(request))
+
+        thread_info = await ThreadsCheckpointerUtil.get_thread_info(thread_id, graph.checkpointer)
+
+        if not allow_access_thread(thread_id, thread_info, request):
+            raise HTTPException(status_code=403, detail="Forbidden")
+
+        try:
+            return thread_info
+        except:
+            raise HTTPException(status_code=404, detail="Thread not found")
+
+    router.include_router(thread_router, prefix="/{thread_id}")
+
+    return router

veri_agents_api/fastapi/threads/test/__main__.py

@@ -0,0 +1,58 @@
+from typing import Annotated
+
+from langchain_aws import ChatBedrock
+from typing_extensions import TypedDict
+
+from fastapi import FastAPI
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+import uvicorn
+
+from veri_agents_api.fastapi.threads import create_threads_router
+
+if __name__ == "__main__":
+    class State(TypedDict):
+        messages: Annotated[list, add_messages]
+
+    graph_builder = StateGraph(State)
+
+    # noinspection PyArgumentList
+    llm = ChatBedrock(model_id="us.anthropic.claude-3-7-sonnet-20250219-v1:0") # pyright: ignore[reportCallIssue]
+
+    def chatbot(state: State):
+        return {"messages": [llm.invoke(state["messages"])]}
+
+    # The first argument is the unique node name
+    # The second argument is the function or object that will be called whenever
+    # the node is used.
+    graph_builder.add_node("chatbot", chatbot)
+    graph_builder.set_entry_point("chatbot")
+    graph_builder.set_finish_point("chatbot")
+
+
+
+    # in-memory persistence
+    memory = MemorySaver()
+    graph = graph_builder.compile(checkpointer=memory)
+
+    # veri-agents convenience router
+    threads_router = create_threads_router(
+        # same graph for every request
+        get_graph=lambda req: graph
+    )
+
+    # root fastapi app
+    app = FastAPI()
+    app.include_router(threads_router)
+
+    uvicorn.run(app, port=5000, log_level="info")
+    # you can now access:
+    #   GET  /openapi.json
+    #   GET  /threads
+    #   GET /threads/{thread_id}/info
+    #   POST /threads/{thread_id}/invoke
+    #   POST /threads/{thread_id}/stream
+    #   GET  /threads/{thread_id}/history
+    #   GET  /threads/{thread_id}/feedback
+    #   POST /threads/{thread_id}/feedback

veri_agents_api/threads_util/__init__.py

@@ -0,0 +1,2 @@
+from .checkpointer import *
+from .schema import *

veri_agents_api/threads_util/checkpointer.py

@@ -0,0 +1,43 @@
+from typing import cast
+
+from .schema import ThreadInfo
+from langgraph.types import Checkpointer
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+class ThreadsCheckpointerUtil:
+    @staticmethod
+    async def get_thread_info(thread_id: str, checkpointer: Checkpointer) -> ThreadInfo | None:
+        if not isinstance(checkpointer, BaseCheckpointSaver):
+            raise Exception("checkpointer must be instance of BaseCheckpointSaver")
+
+        chk_tuple = (await checkpointer.aget_tuple(config={
+            "configurable": {"_has_threadinfo": True, "thread_id": thread_id}}))
+
+        if chk_tuple is None:
+            return None
+
+        thread_metadata = chk_tuple.metadata
+        return ThreadInfo(
+            thread_id=thread_id,
+        )
+
+    @staticmethod
+    async def list_threads(checkpointer: Checkpointer) -> list[ThreadInfo]:
+        if not isinstance(checkpointer, BaseCheckpointSaver):
+            raise Exception("checkpointer must be instance of BaseCheckpointSaver")
+
+        init_step_checkpoints = checkpointer.alist(config=None, filter={ 'step': -1 }) # get initial steps only - this ensures we are only getting one thread_id from the checkpoints of a thread
+
+        all_accessible_thread_ids: list[str] = []
+        async for checkpoint in init_step_checkpoints:
+            thread_id = cast(str | None, checkpoint.config.get("configurable", {"thread_id": None}).get("thread_id", None))
+            if thread_id is not None: # and allow_access_thread(thread_id, request) ?
+                all_accessible_thread_ids.append(cast(str, thread_id))
+
+        all_thread_info: list[ThreadInfo] = []
+        for thread_id in all_accessible_thread_ids:
+            thread_info = await ThreadsCheckpointerUtil.get_thread_info(thread_id, checkpointer)
+            if thread_info is not None:
+                all_thread_info.append(thread_info)
+
+        return all_thread_info

veri_agents_api/threads_util/schema.py

@@ -0,0 +1,13 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+
+class ThreadInfo(BaseModel):
+    """Information about a single thread."""
+
+    thread_id: str
+    # workflow_id: str
+    # name: str
+    # user: str
+    # metadata: dict = Field(default={})
+    # creation: datetime = Field(default_factory=datetime.now)

veri_agents_api/util/awaitable.py

@@ -0,0 +1,11 @@
+import asyncio
+from typing import TypeVar, Awaitable, cast
+
+T = TypeVar('T')
+
+async def as_awaitable(maybe_coroutine: T | Awaitable[T]) -> T:
+    if asyncio.iscoroutine(maybe_coroutine):
+        return await maybe_coroutine
+    return cast(T, maybe_coroutine)
+
+type MaybeAwaitable[T] = T | Awaitable[T]

veri_agents_api-0.1.1.dist-info/METADATA

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: veri-agents-api
+Version: 0.1.1
+Summary: Add your description here
+Author-email: Markus Toman <mtoman@veritone.com>, Teo Boley <tboley@veritone.com>
+Requires-Python: >=3.12
+Requires-Dist: veri-agents-common[langgraph]==0.1.1
+Provides-Extra: dev
+Requires-Dist: langchain-aws>=0.2.21; extra == 'dev'
+Requires-Dist: uvicorn>=0.34.2; extra == 'dev'
+Requires-Dist: veri-agents-common[fastapi,langfuse]==0.1.1; extra == 'dev'
+Provides-Extra: fastapi
+Requires-Dist: veri-agents-common[fastapi,langfuse]==0.1.1; extra == 'fastapi'
+Provides-Extra: fastapi-dev
+Requires-Dist: langchain-aws>=0.2.21; extra == 'fastapi-dev'
+Requires-Dist: uvicorn>=0.34.2; extra == 'fastapi-dev'
+Requires-Dist: veri-agents-common[fastapi,langfuse]==0.1.1; extra == 'fastapi-dev'

veri_agents_api-0.1.1.dist-info/RECORD

@@ -0,0 +1,18 @@
+veri_agents_api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+veri_agents_api/fastapi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+veri_agents_api/fastapi/thread/__init__.py,sha256=WFpfiIPAMaAbD2OFDUlIJ-tt1-bXOz2fOAGRYvyCFts,44
+veri_agents_api/fastapi/thread/router.py,sha256=DDEjb2-xGX3Osg9BBP2AwHOkt7JzO-tL58rWKnmsjME,14135
+veri_agents_api/fastapi/thread/schema.py,sha256=2IQlKuzWzLPeynIHVKKKXhCw3pVZe0KjJkQo1CM8B3Y,5519
+veri_agents_api/fastapi/thread/test/multi_thread.py,sha256=xxe6SpfOLbWeTFItNAX6ipHz2zkOdgtE4P_ADCgWNF4,1888
+veri_agents_api/fastapi/thread/test/single_thread.py,sha256=lbSehG3oWo-bC5iJCVWylkV_TtEdemo87PFappb0Et4,1655
+veri_agents_api/fastapi/threads/__init__.py,sha256=GAq8qINP3yNKm8tnrRGiWuH5IuEsPI4Rzl5T-5Ql1bY,22
+veri_agents_api/fastapi/threads/router.py,sha256=wlYYFHQqJvXB2MZUWF_7xvZE8oTJnj6C9GIi-BWPmhI,2887
+veri_agents_api/fastapi/threads/test/__main__.py,sha256=2Cp7dtrKRBsC_DQM7iim8lF7lOgQk6T0s7BNH56b40A,1770
+veri_agents_api/threads_util/__init__.py,sha256=P06UGWyMeigeiE6dDmsS40k28JZs0NmIOSnldsB7bQ4,50
+veri_agents_api/threads_util/checkpointer.py,sha256=4eiR5zfbjpPqPNzQ7_K8tiFWyWXYxwfxWw6k-fwI774,1904
+veri_agents_api/threads_util/schema.py,sha256=f3l4l5MozS4QvxOewwVSRCKScQmw3Xi2IIpCfbqtsAM,323
+veri_agents_api/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+veri_agents_api/util/awaitable.py,sha256=FerfQ0QKK-nUo-jQ9MfxM8g5pbo0P10Tt1mYZZa1zTs,303
+veri_agents_api-0.1.1.dist-info/METADATA,sha256=McgIAMbZ9e_DtLQUbQYo8M-jkKDntJrIT15c_qFDNWU,777
+veri_agents_api-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+veri_agents_api-0.1.1.dist-info/RECORD,,

veri_agents_api-0.1.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any