session.py

Overview

The session.py file defines classes and methods for managing interactive conversational sessions within the InfiniFlow system. It primarily provides an abstraction for a Session entity that supports multi-turn question-answer interactions with either a "chat" or "agent" backend model. The file encapsulates the logic for sending user queries, streaming or synchronous response handling, and managing session state and messages. Additionally, it defines a Message class representing individual assistant messages returned during the session.

This module acts as a client-side interface for interacting with remote chat or agent completion APIs, handling the network requests and response parsing transparently. It is designed to be extendable and integrates with a base class (Base) which likely provides HTTP request utilities and shared functionality.

Classes and Methods

Class: Session

Represents a conversational session with either a chat or agent backend. Supports sending questions, receiving responses (streamed or complete), and updating session metadata.

Constructor: __init__(self, rag, res_dict)

Initializes a new Session instance.

• Parameters:

  • rag: An instance of a resource or API gateway class used for HTTP communication (inherited functionality).

  • res_dict (dict): A dictionary containing session-related data. Expected keys include:

    • "chat_id" (optional): If present and not None, session is type "chat".

    • "agent_id" (optional): If present and not None, session is type "agent".

• Behavior:

  • Sets default properties: id=None, name="New session".

  • Initializes messages with a default assistant greeting message.

  • Determines session type ("chat" or "agent") based on keys in res_dict.

  • Calls superclass initializer with rag and res_dict.

• Example:

  session = Session(rag_instance, {"chat_id": "1234"})
  

Method: ask(self, question="", stream=True, **kwargs)

Sends a question to the session backend and yields or returns the assistant's responses.

• Parameters:

  • question (str): The user's query to ask. Defaults to empty string.

  • stream (bool): If True, yields response messages as they arrive; if False, returns the complete response message.

  • **kwargs: Additional optional parameters forwarded to the chat backend request.

• Returns:

  • If stream=True: A generator yielding Message objects as the response progresses.

  • If stream=False: A single Message object containing the complete answer.

• Raises:

  • Exception if the response contains an error message or invalid JSON.

• Usage:

  # Streaming usage
  for message in session.ask("Hello, assistant!"):
      print(message.content)
  
  # Non-streaming usage
  response = session.ask("What's the weather?", stream=False)
  print(response.content)
  

• Implementation details:

  • Chooses which internal ask method to call based on session type (_ask_chat or _ask_agent).

  • Handles line-by-line streaming JSON data parsing.

  • Filters out control or status messages.

  • Parses JSON data and constructs Message objects via _structure_answer.

Method: _structure_answer(self, json_data)

Internal helper to parse JSON response data and create a Message instance.

• Parameters:

  • json_data (dict): Parsed JSON dictionary containing assistant answer and optional reference info.

• Returns:

  • A Message object representing the assistant's answer, including any reference chunks if available.

• Details:

  • Extracts "answer" text and optional "reference" chunks.

  • Constructs a dictionary with content, role (set to "assistant"), and reference.

  • Instantiates and returns a Message with this data.

Method: _ask_chat(self, question: str, stream: bool, **kwargs)

Sends a question to the chat backend API.

• Parameters:

  • question (str): The user question.

  • stream (bool): Whether to stream the response.

  • **kwargs: Additional parameters sent in the request JSON.

• Returns:

  • The HTTP response object from the post request.

• Details:

  • Constructs JSON payload with question, stream flag, session ID, plus extra kwargs.

  • Calls the inherited post() method to send request to /chats/{chat_id}/completions.

Method: _ask_agent(self, question: str, stream: bool)

Sends a question to the agent backend API.

• Parameters:

  • question (str): The user question.

  • stream (bool): Whether to stream the response.

• Returns:

  • The HTTP response object from the post request.

• Details:

  • Sends JSON payload with question, stream flag, session ID.

  • Calls post() on /agents/{agent_id}/completions.

Method: update(self, update_message)

Updates the session metadata on the backend.

• Parameters:

  • update_message (dict): Data to update the session with.

• Raises:

  • Exception if the backend response code is not 0, with the error message.

• Details:

  • Sends a PUT request to /chats/{chat_id}/sessions/{id} with update_message.

  • Parses JSON response and validates success.

Class: Message

Represents an individual assistant message in the session conversation.

Constructor: __init__(self, rag, res_dict)

Initializes a new Message.

• Parameters:

  • rag: API gateway or resource handler instance.

  • res_dict (dict): Dictionary containing message data, including:

    • content (str): The message content.

    • reference (optional): Reference chunks related to the message.

    • role (str): Role of the message sender (default "assistant").

    • prompt (optional): The prompt text that generated this message.

    • id (optional): Message identifier.

• Defaults:

  • content defaults to a greeting message.

  • reference defaults to None.

  • role defaults to "assistant".

Implementation Details and Algorithms

• Session Type Detection:
  The constructor inspects res_dict for keys "chat_id" or "agent_id" to determine session type, affecting which API endpoint is used for queries.

• Streaming Response Handling:
  The ask() method supports streaming by iterating over lines of the HTTP response. It decodes and parses JSON messages prefixed with "data:", ignoring heartbeat or running status signals. Error messages embedded in the stream trigger exceptions immediately.

• Response Structuring:
  Responses are parsed into structured Message objects, encapsulating the assistant's textual answer and any reference data (e.g., document chunks cited).

• HTTP Request Abstraction:
  post() and put() methods are inherited from Base and used to communicate with RESTful endpoints, abstracting away HTTP details.

Interaction with Other System Components

• Base Class (Base):
  Session and Message inherit from Base, which likely provides shared attributes and HTTP methods (post, put, etc.), enabling REST API communication.

• RAG (Resource API Gateway):
  The rag parameter passed to constructors is used internally for making HTTP requests and managing resource configurations.

• REST API Endpoints:

  • /chats/{chat_id}/completions — For chat session completions.

  • /agents/{agent_id}/completions — For agent session completions.

  • /chats/{chat_id}/sessions/{id} — For updating session metadata.

• Message Passing:
  The Session creates and yields Message objects representing assistant responses, which can be consumed by UI components or other system modules.

Visual Diagram

classDiagram
    class Session {
        +id: str
        +name: str
        +messages: list
        +ask(question: str, stream: bool, **kwargs) Message | generator
        +update(update_message: dict) void
        -_structure_answer(json_data: dict) Message
        -_ask_chat(question: str, stream: bool, **kwargs) Response
        -_ask_agent(question: str, stream: bool) Response
    }
    class Message {
        +content: str
        +reference: dict | None
        +role: str
        +prompt: str | None
        +id: str | None
    }
    class Base {
        <<abstract>>
        +post(endpoint: str, data: dict, stream: bool) Response
        +put(endpoint: str, data: dict) Response
    }

    Session --|> Base
    Message --|> Base

Summary

The session.py file provides a robust interface for managing conversational sessions with chat or agent AI backends. It abstracts API communication, supports real-time streaming of assistant responses, and encapsulates message data consistently. This module is fundamental for building interactive, turn-based assistant features within the InfiniFlow application, enabling flexible session handling and extensibility.

End of Documentation for session.py