llama-stack-mirror/llama_stack/apis/preprocessing/preprocessing.py
2025-03-03 20:31:06 +01:00

52 lines
1.5 KiB
Python

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from enum import Enum
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
from pydantic import BaseModel
from llama_stack.apis.common.content_types import URL, InterleavedContent
from llama_stack.apis.preprocessing.preprocessors import Preprocessor
from llama_stack.apis.vector_io import Chunk
from llama_stack.schema_utils import json_schema_type, webmethod
class PreprocessingInputType(Enum):
document_content = "document_content"
document_path = "document_path"
@json_schema_type
class PreprocessingInput(BaseModel):
preprocessor_input_id: str
preprocessor_input_type: Optional[PreprocessingInputType]
path_or_content: str | URL
PreprocessorOptions = Dict[str, Any]
@json_schema_type
class PreprocessingResponse(BaseModel):
status: bool
results: Optional[List[str | InterleavedContent | Chunk]]
class PreprocessorStore(Protocol):
def get_preprocessor(self, preprocessor_id: str) -> Preprocessor: ...
@runtime_checkable
class Preprocessing(Protocol):
preprocessor_store: PreprocessorStore
@webmethod(route="/preprocess", method="POST")
async def preprocess(
self,
preprocessor_id: str,
preprocessor_inputs: List[PreprocessingInput],
options: PreprocessorOptions,
) -> PreprocessingResponse: ...