chore: refactoring API and adding stubs for new OpenAI endpoints (#437)

* moving over assistants stuff

* adding in some files endpoints stuff

* Addressing formatting issues

* add stubs for vector store

* add threads stubs

* stub messages and vector store files

* fix openai typing refs

* stubbing out runs and runs steps

* typing issues, and removing tests

* only include endpoints if Supabase is available

* fix formatting issue

* more robustly handle missing supabase

* refactor supabase_wrapper

* snazzy tests

* refactor Supabase Wrapper

* remove supabase-jwt secret (for now)

* update stubs to async defs

* new branch to focus on refactoring API and adding new stubs for OpenAI endpoints

* revert Zarf.yaml changes

* we don't need .env until we have Supabase

* refactoring tests

* fix formatting issues

* removing Supabase stuff (for now)

* refactoring api

* fix formatting issues pre-commit didn't catch...

* add back legacy completions endpoint (for now)

* some rework

* ensure pytest runs

* adding back the custom Usage class

* fixing blank space

* moving types to backend

* updating typing for helpers

* deprecate legacy Optional typing

* deprecate legacy List typing

* fix nonetype error

* fixing Usage convention

.gitignore

@@ -14,6 +14,7 @@ zarf-sbom/
 /values.yaml
 .pytest_cache
 *egg-info
+egg-info/
 build/
 *.whl
 .model/

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
 
   # Python linting and formatting
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.2.2
+    rev: v0.3.4
     hooks:
       - id: ruff # Run the linter.
       - id: ruff-format # Run the formatter.

src/leapfrogai_api/Makefile

@@ -0,0 +1,6 @@
+install:
+	python -m pip install -e .
+
+dev:
+	make install
+	python -m uvicorn main:app --port 3000 --reload

src/leapfrogai_api/backends/openai/grpc_client.py → src/leapfrogai_api/backend/grpc_client.py

@@ -1,10 +1,11 @@
-from typing import Iterator
+"""gRPC client for OpenAI models."""
 
+from typing import Iterator
 import grpc
-import leapfrogai_sdk as lfai
 from fastapi.responses import StreamingResponse
-from leapfrogai_api.backends.openai.helpers import recv_chat, recv_completion
-from leapfrogai_api.backends.openai.types import (
+import leapfrogai_sdk as lfai
+from leapfrogai_api.backend.helpers import recv_chat, recv_completion
+from leapfrogai_api.backend.types import (
     ChatChoice,
     ChatCompletionResponse,
     ChatMessage,
@@ -19,6 +20,7 @@
 
 
 async def stream_completion(model: Model, request: lfai.CompletionRequest):
+    """Stream completion using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.CompletionStreamServiceStub(channel)
         stream = stub.CompleteStream(request)
@@ -31,6 +33,7 @@ async def stream_completion(model: Model, request: lfai.CompletionRequest):
 
 # TODO: Clean up completion() and stream_completion() to reduce code duplication
 async def completion(model: Model, request: lfai.CompletionRequest):
+    """Complete using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.CompletionServiceStub(channel)
         response: lfai.CompletionResponse = await stub.Complete(request)
@@ -50,6 +53,7 @@ async def completion(model: Model, request: lfai.CompletionRequest):
 
 
 async def stream_chat_completion(model: Model, request: lfai.ChatCompletionRequest):
+    """Stream chat completion using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.ChatCompletionStreamServiceStub(channel)
         stream = stub.ChatCompleteStream(request)
@@ -60,6 +64,7 @@ async def stream_chat_completion(model: Model, request: lfai.ChatCompletionReque
 
 # TODO: Clean up completion() and stream_completion() to reduce code duplication
 async def chat_completion(model: Model, request: lfai.ChatCompletionRequest):
+    """Complete chat using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.ChatCompletionServiceStub(channel)
         response: lfai.ChatCompletionResponse = await stub.ChatComplete(request)
@@ -82,6 +87,7 @@ async def chat_completion(model: Model, request: lfai.ChatCompletionRequest):
 
 
 async def create_embeddings(model: Model, request: lfai.EmbeddingRequest):
+    """Create embeddings using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.EmbeddingsServiceStub(channel)
         e: lfai.EmbeddingResponse = await stub.CreateEmbedding(request)
@@ -98,6 +104,7 @@ async def create_embeddings(model: Model, request: lfai.EmbeddingRequest):
 
 
 async def create_transcription(model: Model, request: Iterator[lfai.AudioRequest]):
+    """Transcribe audio using the specified model."""
     async with grpc.aio.insecure_channel(model.backend) as channel:
         stub = lfai.AudioStub(channel)
         response: lfai.AudioResponse = await stub.Transcribe(request)

src/leapfrogai_api/backends/openai/helpers.py → src/leapfrogai_api/backend/helpers.py

@@ -1,8 +1,9 @@
-from typing import BinaryIO, Iterator, Union
+"""Helper functions for the OpenAI backend."""
 
+from typing import BinaryIO, Iterator
 import grpc
 import leapfrogai_sdk as lfai
-from leapfrogai_api.backends.openai.types import (
+from leapfrogai_api.backend.types import (
     ChatCompletionResponse,
     ChatDelta,
     ChatStreamChoice,
@@ -44,6 +45,7 @@ async def recv_chat(
         lfai.ChatCompletionRequest, lfai.ChatCompletionResponse
     ],
 ):
+    """Generator that yields chat completion responses as Server-Sent Events."""
     async for c in stream:
         yield (
             "data: "
@@ -69,7 +71,8 @@ async def recv_chat(
     yield "data: [DONE]\n\n"
 
 
-def grpc_chat_role(role: str) -> Union[lfai.ChatRole, None]:
+def grpc_chat_role(role: str) -> lfai.ChatRole:
+    """Converts a string to a ChatRole."""
     match role:
         case "user":
             return lfai.ChatRole.USER  # type: ignore
@@ -85,6 +88,7 @@ def grpc_chat_role(role: str) -> Union[lfai.ChatRole, None]:
 
 # read_chunks is a helper method that chunks the bytes of a file (audio file) into a iterator of AudioRequests
 def read_chunks(file: BinaryIO, chunk_size: int) -> Iterator[lfai.AudioRequest]:
+    """Reads a file in chunks and yields AudioRequests."""
     while True:
         chunk = file.read(chunk_size)
         if not chunk:

src/leapfrogai_api/backends/openai/types.py → src/leapfrogai_api/backend/types.py

@@ -1,24 +1,52 @@
-from __future__ import annotations
-
-from typing import Dict, Optional
+"""Typing definitions for assistants API."""
 
-from fastapi import File, Form, UploadFile
+from __future__ import annotations
+from typing import Literal
 from pydantic import BaseModel
-
+from fastapi import UploadFile, Form, File
 
 ##########
 # GENERIC
 ##########
+
+
 class Usage(BaseModel):
+    """Usage object."""
+
     prompt_tokens: int
     completion_tokens: int | None = None
     total_tokens: int
 
 
 ##########
-# COMPLETION
+# MODELS
 ##########
+
+
+class ModelResponseModel(BaseModel):
+    """Response object for models."""
+
+    id: str
+    object: str = "model"
+    created: int = 0
+    owned_by: str = "leapfrogai"
+
+
+class ModelResponse(BaseModel):
+    """Response object for models."""
+
+    object: str = "list"
+    data: list[ModelResponseModel] = []
+
+
+############
+# COMPLETION
+############
+
+
 class CompletionRequest(BaseModel):
+    """Request object for completion."""
+
     model: str
     prompt: str | list[int]
     stream: bool | None = False
@@ -27,13 +55,17 @@ class CompletionRequest(BaseModel):
 
 
 class CompletionChoice(BaseModel):
+    """Choice object for completion."""
+
     index: int
     text: str
     logprobs: object = None
     finish_reason: str = ""
 
 
 class CompletionResponse(BaseModel):
+    """Response object for completion."""
+
     id: str = ""
     object: str = "completion"
     created: int = 0
@@ -48,22 +80,30 @@ class CompletionResponse(BaseModel):
 
 
 class ChatFunction(BaseModel):
+    """Function object for chat completion."""
+
     name: str
-    parameters: Dict[str, object]
+    parameters: dict[str, object]
     description: str
 
 
 class ChatMessage(BaseModel):
+    """Message object for chat completion."""
+
     role: str
     content: str
 
 
 class ChatDelta(BaseModel):
+    """Delta object for chat completion."""
+
     role: str
     content: str | None = ""
 
 
 class ChatCompletionRequest(BaseModel):
+    """Request object for chat completion."""
+
     model: str
     messages: list[ChatMessage]
     functions: list | None = None
@@ -75,19 +115,23 @@ class ChatCompletionRequest(BaseModel):
 
 
 class ChatChoice(BaseModel):
+    """Choice object for chat completion."""
+
     index: int
     message: ChatMessage
     finish_reason: str | None = ""
 
 
 class ChatStreamChoice(BaseModel):
+    """Stream choice object for chat completion."""
+
     index: int
     delta: ChatDelta
     finish_reason: str | None = ""
 
 
 class ChatCompletionResponse(BaseModel):
-    """https://platform.openai.com/docs/api-reference/chat/object"""
+    """Response object for chat completion."""
 
     id: str = ""
     object: str = "chat.completion"
@@ -97,36 +141,34 @@ class ChatCompletionResponse(BaseModel):
     usage: Usage | None = None
 
 
+#############
+# EMBEDDINGS
+#############
+
+
 class CreateEmbeddingRequest(BaseModel):
+    """Request object for creating embeddings."""
+
     model: str
     input: str | list[str] | list[int] | list[list[int]]
     user: str | None = None
 
 
 class EmbeddingResponseData(BaseModel):
+    """Response object for embeddings."""
+
     embedding: list[float]
     index: int
     object: str = "embedding"
 
 
 class CreateEmbeddingResponse(BaseModel):
+    """Response object for embeddings."""
+
     data: list[EmbeddingResponseData]
     model: str
     object: str = "list"
-    usage: Usage
-
-
-# yes I know, this is a pure API response class for matching OpenAI
-class ModelResponseModel(BaseModel):
-    id: str
-    object: str = "model"
-    created: int = 0
-    owned_by: str = "leapfrogai"
-
-
-class ModelResponse(BaseModel):
-    object: str = "list"
-    data: list[ModelResponseModel] = []
+    usage: Usage | None = None
 
 
 ##########
@@ -140,17 +182,17 @@ class CreateTranscriptionRequest(BaseModel):
     language: str = ""
     prompt: str = ""
     response_format: str = ""
-    temperature: float = 1
+    temperature: float = 1.0
 
     @classmethod
     def as_form(
         cls,
         file: UploadFile = File(...),
         model: str = Form(...),
-        language: Optional[str] = Form(""),
-        prompt: Optional[str] = Form(""),
-        response_format: Optional[str] = Form(""),
-        temperature: Optional[float] = Form(1),
+        language: str | None = Form(""),
+        prompt: str | None = Form(""),
+        response_format: str | None = Form(""),
+        temperature: float | None = Form(1.0),
     ) -> CreateTranscriptionRequest:
         return cls(
             file=file,
@@ -164,3 +206,50 @@ def as_form(
 
 class CreateTranscriptionResponse(BaseModel):
     text: str
+
+
+#############
+# FILES
+#############
+
+
+class UploadFileRequest(BaseModel):
+    """Request object for uploading a file."""
+
+    file: UploadFile
+    purpose: Literal["assistants"] | None = "assistants"
+
+    @classmethod
+    def as_form(
+        cls,
+        file: UploadFile = File(...),
+        purpose: str | None = Form("assistants"),
+    ) -> UploadFileRequest:
+        """Create an instance of the class from form data."""
+        return cls(file=file, purpose=purpose)
+
+
+#############
+# ASSISTANTS
+#############
+
+
+class CreateAssistantRequest(BaseModel):
+    """Request object for creating an assistant."""
+
+    model: str = "mistral"
+    name: str | None = "Froggy Assistant"
+    description: str | None = "A helpful assistant."
+    instructions: str | None = "You are a helpful assistant."
+    tools: list[dict[Literal["type"], Literal["file_search"]]] | None = [
+        {"type": "file_search"}
+    ]  # This is all we support right now
+    tool_resources: object | None = {}
+    metadata: object | None = {}
+    temperature: float | None = 1.0
+    top_p: float | None = 1.0
+    response_format: Literal["auto"] | None = "auto"  # This is all we support right now
+
+
+class ModifyAssistantRequest(CreateAssistantRequest):
+    """Request object for modifying an assistant."""