DeepWiki WebSocket API Implementation for Real-Time Chat Streaming

The DeepWiki WebSocket API streams real-time chat responses through an async generator in api/websocket_wiki.py, accepting JSON requests, initializing a per-request RAG pipeline, and iterating over provider-specific streaming chunks to deliver incremental text updates.

The DeepWiki WebSocket API enables low-latency, bidirectional communication for AI-powered chat completions in the AsyncFuncAI/deepwiki-open repository. This implementation handles connection lifecycle management, retrieval-augmented generation (RAG) context injection, and multi-provider streaming through a single endpoint.

WebSocket Endpoint Architecture

The core implementation resides in api/websocket_wiki.py, specifically within the handle_websocket_chat function. This async handler manages the entire lifecycle of a WebSocket connection, from initial handshake to final disconnection, while coordinating between the client, the RAG retrieval system, and various large language model (LLM) providers.

Connection Lifecycle and Request Handling

Accepting the Connection and Receiving Payloads

The handler begins by establishing the bidirectional channel:

await websocket.accept()

Once connected, the server awaits a JSON payload containing the chat completion parameters:

request = ChatCompletionRequest(**await websocket.receive_json())

This payload is validated against the ChatCompletionRequest Pydantic model, ensuring type safety for fields like provider, model, and messages.

Token Count Guardrails

Before processing, the implementation checks prompt size to prevent context window overflow:

token_count = count_tokens(request.messages)
if token_count > 8000:
    logger.warning(f"Large prompt detected: {token_count} tokens")

The count_tokens utility analyzes the message list and logs warnings when prompts exceed approximately 8,000 tokens, allowing clients to adjust their requests before hitting provider-specific limits.

RAG Pipeline Initialization

For each incoming request, the system instantiates a fresh RAG object to handle context retrieval:

rag = RAG(provider=request.provider, model=request.model)

The retriever is configured with repository-specific filters to narrow the search space:

retriever = rag.get_retriever(
    repo_name=request.repo_name,
    filters={"source": request.source_type}
)

This per-request isolation ensures that conversation context remains separate across concurrent connections, while the retrieval system fetches relevant documentation chunks based on the current query.

Message Validation and Prompt Construction

Validating Conversation State

The handler enforces structural requirements on the message list:

if not request.messages:
    raise ValueError("At least one user message required")
    
if request.messages[-1].role != "user":
    raise ValueError("Last message must be from user")

These checks ensure that the LLM receives a valid conversation thread ending with a user query, preventing malformed requests from reaching the provider APIs.

Building the Contextual Prompt

The final prompt aggregates multiple context sources:

prompt_parts = [
    system_prompt,
    conversation_history,
    file_content if request.file_path else "",
    retrieved_context
]
full_prompt = "\n\n".join(filter(None, prompt_parts))

This construction concatenates the system instructions, previous conversation turns, optional file attachments, and RAG-retrieved documentation into a single coherent context string for the LLM.

Provider Selection and Streaming Implementation

Branching to Provider-Specific Clients

The implementation supports seven distinct LLM providers through a branching structure in api/websocket_wiki.py:

if request.provider == "ollama":
    client = OllamaClient()
elif request.provider == "openrouter":
    client = OpenRouterClient()
elif request.provider == "openai":
    client = OpenAIClient()

# ... additional providers: bedrock, azure, dashscope, google

Each client is initialized with model-specific kwargs retrieved from api/config.py, ensuring proper authentication and endpoint configuration.

Converting Inputs and Initiating Streams

Before calling the model, inputs are converted to provider-specific API arguments:

api_kwargs = client.convert_inputs_to_api_kwargs(
    prompt=full_prompt,
    model=request.model,
    temperature=request.temperature
)

The acall method returns an async generator that yields response chunks:

stream = await model.acall(
    api_kwargs=api_kwargs,
    model_type=ModelType.LLM
)

async for chunk in stream:
    text = extract_text(chunk)
    await websocket.send_text(text)

The handler iterates over this generator, extracts the textual content from each chunk, strips internal markers, and immediately forwards the text to the client via websocket.send_text(), achieving true real-time streaming.

Summary

• The DeepWiki WebSocket API implementation centers on api/websocket_wiki.py, specifically the handle_websocket_chat function.
• Each connection initializes an isolated RAG pipeline to retrieve contextually relevant documentation before generation begins.
• The system enforces token limits and message validation to prevent malformed requests from reaching LLM providers.
• Seven provider backends (Ollama, OpenRouter, OpenAI, Bedrock, Azure, DashScope, Google) are supported through a unified async streaming interface.
• Responses stream in real-time via an async generator pattern, with chunks immediately forwarded to the client over the persistent WebSocket connection.

Frequently Asked Questions

What file contains the main WebSocket handler for DeepWiki?

The primary implementation resides in api/websocket_wiki.py, specifically within the handle_websocket_chat async function that manages the full connection lifecycle from handshake to disconnection.

How does DeepWiki handle context retrieval for chat requests?

Each incoming request instantiates a fresh RAG object configured with repository-specific filters. The retriever fetches relevant documentation chunks based on the current query and repository name, injecting this context into the final prompt sent to the LLM.

Which LLM providers are supported by the WebSocket API?

The implementation supports seven providers: Ollama, OpenRouter, OpenAI, AWS Bedrock, Azure OpenAI, DashScope, and Google Generative AI. Each provider has a dedicated client class that handles authentication and API-specific argument conversion.

How is real-time streaming achieved in the WebSocket implementation?

The system uses an async generator pattern where model.acall() returns a stream object. The handler iterates over this generator with async for, extracts text from each chunk, and immediately sends it to the client via websocket.send_text(), ensuring low-latency incremental delivery without waiting for the full response.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}