contentgrid-assistant-api 0.0.2__py3-none-any.whl

contentgrid_assistant_api/routers/message_router.py

@@ -0,0 +1,256 @@
+# type : ignore
+
+import base64
+import json
+from typing import Annotated, AsyncGenerator, List, Union, Optional
+import typing
+import uuid
+from enum import Enum
+from langgraph.graph.state import CompiledStateGraph
+from fastapi import APIRouter, BackgroundTasks, File, Form, Request, UploadFile, status
+from fastapi.params import Depends
+from contentgrid_assistant_api.config import AssistantExtensionConfig
+from fastapi.responses import StreamingResponse
+from contentgrid_assistant_api.db.repositories.thread_repository import ThreadRepository
+from contentgrid_assistant_api.db.types.message import HALHumanMessage, HALAIMessage, HALSystemMessage, HALToolMessage
+from contentgrid_assistant_api.dependencies import DependencyResolver
+from contentgrid_extension_helpers.authentication import ContentGridUser
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from contentgrid_extension_helpers.responses.hal import FastAPIHALCollection, HALTemplateFor, HALLinkFor
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, BaseMessage, ToolMessage
+from fastapi import HTTPException
+
+
+# ContentBlocks
+from langchain_core.messages.content import create_image_block, create_file_block, create_audio_block, create_text_block, ImageContentBlock, AudioContentBlock, FileContentBlock, TextContentBlock
+import logging
+
+from contentgrid_assistant_api.types.context import DefaultThreadContext
+
+def generate_agent_message_router(dep_resolver: DependencyResolver, extension_config: AssistantExtensionConfig, tags: Optional[List[str | Enum]]=None) -> APIRouter:
+    messagesrouter = APIRouter(prefix="/{thread_id}" + extension_config.routes_message_prefix, tags=tags or ["messages"])
+
+    def convert_to_hal_message(message: BaseMessage, thread_id: uuid.UUID) -> Union[HALHumanMessage, HALAIMessage, HALSystemMessage, HALToolMessage]:
+        """Convert a LangChain BaseMessage to the appropriate HAL message type"""
+        if isinstance(message, ToolMessage):
+            # Handle ToolMessage first since it might inherit from other message types
+            hal_message = HALToolMessage(**message.model_dump())
+        elif isinstance(message, HumanMessage):
+            hal_message = HALHumanMessage(**message.model_dump())
+        elif isinstance(message, AIMessage):
+            hal_message = HALAIMessage(**message.model_dump())
+        elif isinstance(message, SystemMessage):
+            hal_message = HALSystemMessage(**message.model_dump())
+        else:
+            # Default to HALSystemMessage for unknown types
+            logging.warning("Unknown type : " + str(type(message)))
+            hal_message = HALSystemMessage(**message.model_dump())
+        
+        # Add HAL links and templates
+        hal_message.links = {
+            "self": HALLinkFor(endpoint_function_name="read_message", tags=tags, path_params={"thread_id": str(thread_id), "message_id": str(message.id)}, condition=message.id is not None),
+            "thread": HALLinkFor(endpoint_function_name="read_thread", tags=tags, path_params={"thread_id": str(thread_id)}),
+            "messages": HALLinkFor(endpoint_function_name="read_messages", tags=tags, path_params={"thread_id": str(thread_id)})
+        }
+        
+        return hal_message
+
+    async def convert_upload_file_to_contentblock(file: UploadFile) -> ImageContentBlock | FileContentBlock | AudioContentBlock:
+        # Reference: https://docs.langchain.com/oss/python/langchain/messages#content-block-reference
+        
+        # Read file content
+        file_content = await file.read()
+        file_base64 = base64.b64encode(file_content).decode("utf-8")
+        
+        # Determine content type based on file's content type
+        if file.content_type:
+            content_type = file.content_type
+            
+            if content_type.startswith("image/"):
+                # Handle image files
+                return create_image_block(
+                        base64=file_base64,
+                        mime_type=content_type
+                    )
+            elif content_type.startswith("audio/"):
+                # Handle audio files
+                return create_audio_block(
+                        base64=file_base64,
+                        mime_type=content_type
+                    )
+            else:
+                # Handle PDF files and other extensions
+                return create_file_block(
+                    base64=file_base64,
+                    mime_type=content_type,
+                    filename=file.filename if file.filename else None
+                )
+        else:
+            # No content type provided, treat as generic file
+            return create_file_block(
+                    base64=file_base64,
+                    mime_type="application/octet-stream",
+                    filename=file.filename if file.filename else None
+                )
+        
+            
+    @messagesrouter.get("/", response_model=FastAPIHALCollection, response_model_exclude_none=True)
+    def read_messages(
+        thread_id: uuid.UUID,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        agent: CompiledStateGraph = Depends(dep_resolver.get_agent_dependency()),
+        thread_context: DefaultThreadContext = Depends(dep_resolver.get_thread_context_dependency())
+    ):
+        """Get all messages for a thread"""
+        state = agent.get_state(config={"configurable" : thread_context}) # type: ignore
+        new_message_can_be_added = True
+        if state and state.values:
+            messages = state.values['messages']
+            # Convert LangChain messages to HAL message types
+            hal_messages = [convert_to_hal_message(msg, thread_id) for msg in messages]
+            last_message = messages[-1]
+            if not isinstance(last_message, AIMessage):
+                new_message_can_be_added = False
+            else:
+                # we have an AI message but we need to check if it is waiting for tools to execute.
+                if len(last_message.tool_calls) > 0:
+                    new_message_can_be_added = False
+        else:
+            hal_messages = []
+            
+        return FastAPIHALCollection[Union[HALHumanMessage, HALAIMessage, HALSystemMessage, HALToolMessage]](
+            _embedded={"messages": hal_messages},
+            _links={
+                "self": HALLinkFor(endpoint_function_name="read_messages", tags=tags, path_params={"thread_id": str(thread_id)}),
+                "thread": HALLinkFor(endpoint_function_name="read_thread", tags=tags, path_params={"thread_id": str(thread_id)})
+            },
+            _templates={"addMessage": HALTemplateFor(endpoint_function_name="add_message", tags=tags, path_params={"thread_id": str(thread_id)}, condition=new_message_can_be_added)}
+        )
+        
+        
+    @messagesrouter.post("/", response_model=HALHumanMessage, status_code=status.HTTP_202_ACCEPTED, response_model_exclude_none=True)
+    async def add_message(
+        request: Request,
+        thread_id: uuid.UUID,
+        question: Annotated[str, Form()],
+        background_tasks : BackgroundTasks,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        agent: CompiledStateGraph = Depends(dep_resolver.get_agent_dependency()),
+        checkpointer: BaseCheckpointSaver = Depends(dep_resolver.get_langgraph_checkpointer_dependency()),
+        thread_context: DefaultThreadContext = Depends(dep_resolver.get_thread_context_dependency()),
+        file: UploadFile = File(None),
+    ):
+        """Add new messages to a thread"""
+            # Check Accept header to determine response format
+        accept_header = request.headers.get("accept", "").lower()
+        streaming = (
+            "text/event-stream" in accept_header
+        )
+        # Build content blocks
+        content_blocks : List[ImageContentBlock | FileContentBlock | AudioContentBlock | TextContentBlock] = [create_text_block(text=question)]
+        
+        if file:
+            # Add custom content blocks based on file type
+            file_block = await convert_upload_file_to_contentblock(file)
+            content_blocks.append(file_block)
+        
+        # Create the message with proper content structure
+        new_message = HumanMessage(content=content_blocks) # type: ignore
+        messages = [new_message]
+        
+        if streaming:
+            state = agent.get_state(config={"configurable" : thread_context}) #type: ignore
+            nb_current_messages = 0
+            if state and state.values and 'messages' in state.values.keys() and len(state.values['messages']):
+                nb_current_messages = len(state.values['messages'])
+
+            @typing.no_type_check
+            async def generate_stream(current_message_index) -> AsyncGenerator[str, None]:
+                for mode, chunk in agent.stream(
+                    {"messages": messages},
+                    {"configurable": thread_context, "recursion_limit": extension_config.graph_recursion_limit},
+                    context=thread_context,
+                    stream_mode=["values", "messages"]
+                ):
+                    if mode=="values":
+                        values = chunk
+                        new_state_messages = values["messages"][current_message_index:]
+                        hal_messages = [convert_to_hal_message(msg, thread_id) for msg in new_state_messages]
+                        
+                        # Convert HAL messages to dictionaries
+                        messages_data = []
+                        for msg in hal_messages:
+                            if hasattr(msg, 'model_dump'):
+                                messages_data.append(msg.model_dump(exclude_unset=True))
+                            else:
+                                messages_data.append(msg)
+                        
+                        # Send JSON chunk with proper structure
+                        yield "event: message\n"
+                        yield f"data: {json.dumps(messages_data)}\n\n"
+                        current_message_index += len(new_state_messages)
+                    elif mode == "messages":
+                        message_chunk, metadata = chunk
+                        if message_chunk.content:
+                            try:
+                                yield "event: token\n"
+                                yield f"data: {json.dumps(message_chunk.content)}\n"
+                                yield f"id: {message_chunk.id}\n\n"
+                            except Exception as e:
+                                logging.exception(e)
+                                continue
+
+                # Send completion signal
+                yield f"data: {json.dumps({'complete': True})}\n\n"
+            
+            return StreamingResponse(
+                generate_stream(current_message_index=nb_current_messages),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                }
+            )
+        else:
+            def process_agent_response():
+                try:
+                    agent.invoke(
+                        {"messages": messages},
+                        {"configurable": thread_context, "recursion_limit": extension_config.graph_recursion_limit},
+                        context=thread_context,
+                    )
+                except Exception as e:
+                    # Log the error since this runs in background
+                    logging.exception(f"Error processing agent response: {e}")
+            
+            # Add the task to background tasks
+            background_tasks.add_task(process_agent_response)
+            
+            # Return immediately with the human message
+            return convert_to_hal_message(new_message, thread_id=thread_id)
+
+    @messagesrouter.get("/{message_id}", response_model=Union[HALHumanMessage, HALAIMessage, HALSystemMessage, HALToolMessage], response_model_exclude_none=True)
+    def read_message(
+        thread_id: uuid.UUID,
+        message_id: str,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        checkpointer: BaseCheckpointSaver = Depends(dep_resolver.get_langgraph_checkpointer_dependency()),
+        thread_context: DefaultThreadContext = Depends(dep_resolver.get_thread_context_dependency())
+    ):
+        """Get a specific message by ID from a thread"""
+        state = checkpointer.get({"configurable": thread_context}) #type: ignore
+        if state and "channel_values" in state.keys():
+            messages = state["channel_values"]["messages"]
+            
+            # Find the message by ID
+            for msg in messages:
+                if msg.id == message_id:
+                    return convert_to_hal_message(msg, thread_id)
+        
+        # Message not found
+        raise HTTPException(status_code=404, detail="Message not found")
+    
+    return messagesrouter

contentgrid_assistant_api/routers/thread_router.py

@@ -0,0 +1,129 @@
+
+from enum import Enum
+from typing import Annotated, List, Optional
+import uuid
+from fastapi import APIRouter, Depends, Query, status
+from pydantic import HttpUrl
+
+from contentgrid_assistant_api.db.repositories.thread_repository import ThreadRepository
+from contentgrid_assistant_api.db.types.thread import ThreadRead, ThreadUpdate
+from contentgrid_assistant_api.dependencies import DependencyResolver
+from contentgrid_extension_helpers.responses.hal import FastAPIHALCollection, HALLinkFor, HALTemplateFor
+from contentgrid_extension_helpers.authentication import ContentGridUser
+from contentgrid_assistant_api.routers.message_router import generate_agent_message_router
+from contentgrid_assistant_api.types.agents import AgentToolCollectionResponse, AgentToolResponse
+from contentgrid_assistant_api.types.context import DefaultThreadContext
+from langchain.messages import HumanMessage
+from langgraph.graph.state import CompiledStateGraph
+from contentgrid_assistant_api.config import AssistantExtensionConfig
+    
+def generate_agent_thread_router(dep_resolver: DependencyResolver, extension_config: AssistantExtensionConfig, tags: Optional[List[str | Enum]]=None):
+    threadrouter = APIRouter(prefix=extension_config.routes_thread_prefix, tags=tags or ["threads"])
+
+    threadrouter.include_router(
+        generate_agent_message_router(dep_resolver, extension_config, tags=tags)
+    )
+
+    @threadrouter.post("/", response_model=ThreadRead, status_code=status.HTTP_201_CREATED, response_model_exclude_none=True)
+    def create_thread(
+        origin: Optional[HttpUrl] = None,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        agent : CompiledStateGraph = Depends(dep_resolver.get_agent_dependency())
+    ):
+        """Create new thread"""
+        
+        thread_id = uuid.uuid4()
+        messages = [HumanMessage(content=extension_config.opening_message)]
+        context = dep_resolver.agent.thread_context(thread_id=str(thread_id), user=user, origin=origin)
+        agent.invoke(
+            {"messages": messages},
+            {"configurable": context}, #type: ignore
+            context=context, #type: ignore
+        )
+        
+        created_thread = thread_repo.create(user=user, origin=origin, component="datamodel", thread_id=thread_id)
+        return ThreadRead(**created_thread.model_dump(), tags=tags)
+
+    @threadrouter.get("/", response_model=FastAPIHALCollection, response_model_exclude_none=True)
+    def read_threads(
+        origin: Optional[HttpUrl] = None,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        offset: int = 0,
+        limit: Annotated[int, Query(le=100)] = 100,
+    ):
+        """Get all threads with pagination"""
+        threads = []
+        if origin:
+            threads = thread_repo.get_all_for_user_and_origin(user, origin, offset=offset, limit=limit)
+        else:
+            threads = thread_repo.get_all_for_user(user, offset=offset, limit=limit)
+        thread_reads = [ThreadRead(**thread.model_dump(), tags=tags) for thread in threads]
+        return FastAPIHALCollection[ThreadRead](
+            _embedded={"threads" : thread_reads}, 
+            _links={"self": HALLinkFor(endpoint_function_name="read_threads", tags=tags, templated=True, params={"offset": offset, "limit": limit})},
+            _templates={"startThread": HALTemplateFor(endpoint_function_name="create_thread", tags=tags, params={"origin": origin.encoded_string()} if origin else {})},
+        )
+
+    @threadrouter.get("/{thread_id}", response_model=ThreadRead, response_model_exclude_none=True)
+    def read_thread(
+        thread_id: uuid.UUID,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency())
+    ):
+        """Get thread by ID"""
+        thread = thread_repo.get_by_id_for_user(thread_id, user)
+        return ThreadRead(**thread.model_dump(), tags=tags)
+
+    @threadrouter.get("/{thread_id}/tools", response_model=AgentToolCollectionResponse, response_model_exclude_none=True)
+    def get_thread_tools(
+        thread_id: uuid.UUID,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency()),
+        thread_context: DefaultThreadContext = Depends(dep_resolver.get_thread_context_dependency())
+    ):
+        """Get available tools for a thread"""
+        tools_list = dep_resolver.agent.get_tools(thread_context)
+        
+        # Extract tool information including schema
+        tool_responses = []
+        for tool in tools_list:
+            tool_dict = tool.model_dump()
+            
+            # Extract schema if args_schema is available
+            if hasattr(tool, 'args_schema') and tool.args_schema:
+                try:
+                    # Get JSON schema from Pydantic model
+                    tool_dict['args'] = tool.args_schema.model_json_schema()
+                except (AttributeError, Exception):
+                    # Fallback if schema doesn't exist or can't be extracted
+                    tool_dict['args'] = {}
+            else:
+                tool_dict['args'] = {}
+            
+            tool_responses.append(AgentToolResponse(**tool_dict))
+        
+        return AgentToolCollectionResponse(thread_id=thread_id, _embedded={"tools": tool_responses}, tags=tags)
+
+    @threadrouter.patch("/{thread_id}", response_model=ThreadRead, response_model_exclude_none=True)
+    def update_thread(
+        thread_id: uuid.UUID,
+        thread_update: ThreadUpdate,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency())
+    ):
+        """Update thread"""
+        updated_thread = thread_repo.update_for_user(thread_id, thread_update, user)
+        return ThreadRead(**updated_thread.model_dump(), tags=tags)
+
+    @threadrouter.delete("/{thread_id}", status_code=status.HTTP_204_NO_CONTENT)
+    def delete_thread(
+        thread_id: uuid.UUID,
+        user: ContentGridUser = Depends(dep_resolver.get_current_user_dependency()),
+        thread_repo: ThreadRepository = Depends(dep_resolver.get_thread_repository_dependency())
+    ):
+        """Delete thread"""
+        thread_repo.delete_for_user(thread_id, user)
+    
+    return threadrouter

contentgrid_assistant_api/types/agents.py

@@ -0,0 +1,50 @@
+
+from typing import Any, Callable, List, Optional
+from enum import Enum
+import uuid
+from pydantic import BaseModel, Field, ConfigDict
+from langchain.tools import BaseTool
+from contentgrid_extension_helpers.responses.hal import FastAPIHALResponse, FastAPIHALCollection, HALLinkFor
+from contentgrid_assistant_api.types.context import DefaultThreadContext
+
+class Agent(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    name : str = "Default agent name"
+    version : str = "v0"
+    get_agent_override : Callable[..., Any] = Field(exclude=True, default_factory=lambda x : None) # used to override the get agent dependency
+    get_current_user_override : Callable[..., Any] = Field(exclude=True, default_factory=lambda x : None) # used to override authentication
+    thread_context : type = Field(exclude=True, default=DefaultThreadContext)
+    # Function that accepts the thread_context and returns the tools.
+    # We do this dynamically because tools could be generated based on the context (like based on the origin profile)
+    get_tools : Callable[..., List[BaseTool]] = Field(exclude=True, default=lambda _ : [])
+    
+    
+class AgentHomeResponse(FastAPIHALResponse):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    name : str
+    version : str
+    
+    def __init__(self, tags: Optional[List[str | Enum]]=None, **kwargs):
+        super().__init__(**kwargs)
+        self.links = {
+            "self" : HALLinkFor(endpoint_function_name="get_agent_home", tags=tags),
+            "threads" : HALLinkFor(endpoint_function_name="read_threads", tags=tags),
+        }
+    
+class AgentToolResponse(FastAPIHALResponse):
+    name : str
+    description : str
+    args : dict
+        
+class AgentToolCollectionResponse(FastAPIHALCollection[AgentToolResponse]):
+    thread_id : uuid.UUID
+    
+    def __init__(self, tags: Optional[List[str | Enum]]=None, **kwargs):
+        super().__init__(**kwargs)
+        self.links = {
+            "self": HALLinkFor(endpoint_function_name="get_thread_tools", tags=tags, templated=False, path_params=lambda instance: {"thread_id": instance.thread_id}),
+            "thread": HALLinkFor(endpoint_function_name="read_thread", tags=tags, templated=False, path_params=lambda instance: {"thread_id": instance.thread_id}),
+        }
+
+    

contentgrid_assistant_api/types/context.py

@@ -0,0 +1,9 @@
+from typing import Optional
+from langchain.agents import AgentState
+from contentgrid_extension_helpers.authentication import ContentGridUser
+
+class DefaultThreadContext(AgentState):
+    user: ContentGridUser
+    thread_id : str
+    origin: Optional[str]
+