phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2026-01-01 08:40:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Initial implementation of RAG operator using the preprocessing endpoint.

Browse source
This commit is contained in:
ilya-kolchinsky 2025-03-05 13:43:26 +01:00
parent c2bd31eb5c
commit 16764a2f06
8 changed files with 74 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

34
llama_stack/providers/inline/preprocessing/basic/basic.py
View file

@ -5,13 +5,14 @@
# the root directory of this source tree.
import logging
import re
from typing import List
from typing import List, Optional
import httpx
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.preprocessing import (
Preprocessing,
PreprocessingDataFormat,
PreprocessingDataType,
PreprocessingInput,
PreprocessingResponse,
@ -54,22 +55,26 @@ class InclineBasicPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate):
self,
preprocessor_id: str,
preprocessor_inputs: List[PreprocessingInput],
options: PreprocessorOptions,
options: Optional[PreprocessorOptions] = None,
) -> PreprocessingResponse:
results = []
for inp in preprocessor_inputs:
is_pdf = options["binary_document_type"] == "pdf"
input_type = self._resolve_input_type(inp, is_pdf)
input_type = self._resolve_input_type(inp)
if input_type == PreprocessingDataType.document_uri:
document = await self._fetch_document(inp, is_pdf)
document = await self._fetch_document(inp)
if document is None:
continue
elif input_type == PreprocessingDataType.binary_document:
document = inp.path_or_content
if not is_pdf:
log.error(f"Unsupported binary document type: {options['binary_document_type']}")
if inp.preprocessor_input_format is None:
log.error(f"Binary document format is not provided for {inp.preprocessor_input_id}, skipping it")
continue
if inp.preprocessor_input_format != PreprocessingDataFormat.pdf:
log.error(
f"Unsupported binary document type {inp.preprocessor_input_format} for {inp.preprocessor_input_id}, skipping it"
)
continue
elif input_type == PreprocessingDataType.raw_text_document:
document = interleaved_content_as_str(inp.path_or_content)
@ -77,7 +82,7 @@ class InclineBasicPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate):
log.error(f"Unexpected preprocessor input type: {inp.preprocessor_input_type}")
continue
if is_pdf:
if inp.preprocessor_input_format == PreprocessingDataFormat.pdf:
document = parse_pdf(document)
results.append(document)
@ -85,7 +90,7 @@ class InclineBasicPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate):
return PreprocessingResponse(status=True, results=results)
@staticmethod
async def _resolve_input_type(preprocessor_input: PreprocessingInput, is_pdf: bool) -> PreprocessingDataType:
async def _resolve_input_type(preprocessor_input: PreprocessingInput) -> PreprocessingDataType:
if preprocessor_input.preprocessor_input_type is not None:
return preprocessor_input.preprocessor_input_type
@ -93,13 +98,13 @@ class InclineBasicPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate):
return PreprocessingDataType.document_uri
if InclineBasicPreprocessorImpl.URL_VALIDATION_PATTERN.match(preprocessor_input.path_or_content):
return PreprocessingDataType.document_uri
if is_pdf:
if preprocessor_input.preprocessor_input_format == PreprocessingDataFormat.pdf:
return PreprocessingDataType.binary_document
return PreprocessingDataType.raw_text_document
@staticmethod
async def _fetch_document(preprocessor_input: PreprocessingInput, is_pdf: bool) -> str | None:
async def _fetch_document(preprocessor_input: PreprocessingInput) -> str | None:
if isinstance(preprocessor_input.path_or_content, str):
url = preprocessor_input.path_or_content
if not InclineBasicPreprocessorImpl.URL_VALIDATION_PATTERN.match(url):
@ -118,4 +123,9 @@ class InclineBasicPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate):
async with httpx.AsyncClient() as client:
r = await client.get(url)
return r.content if is_pdf else r.text
return r.content if preprocessor_input.preprocessor_input_format == PreprocessingDataFormat.pdf else r.text
@staticmethod
def is_pdf(preprocessor_input: PreprocessingInput):
return

4
llama_stack/providers/inline/preprocessing/docling/docling.py
View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import logging
from typing import List
from typing import List, Optional
from docling.document_converter import DocumentConverter
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
@ -51,7 +51,7 @@ class InclineDoclingPreprocessorImpl(Preprocessing, PreprocessorsProtocolPrivate
self,
preprocessor_id: str,
preprocessor_inputs: List[PreprocessingInput],
options: PreprocessorOptions,
options: Optional[PreprocessorOptions] = None,
) -> PreprocessingResponse:
results = []

10
llama_stack/providers/inline/preprocessing/simple_chunking/simple_chunking.py
View file

@ -5,7 +5,7 @@
# the root directory of this source tree.
import logging
from enum import Enum
from typing import List, Tuple
from typing import List, Optional, Tuple
from llama_models.llama3.api import Tokenizer
@ -49,7 +49,7 @@ class InclineSimpleChunkingImpl(Preprocessing, PreprocessorsProtocolPrivate):
self,
preprocessor_id: str,
preprocessor_inputs: List[PreprocessingInput],
options: PreprocessorOptions,
options: Optional[PreprocessorOptions] = None,
) -> PreprocessingResponse:
chunks = []
@ -64,9 +64,11 @@ class InclineSimpleChunkingImpl(Preprocessing, PreprocessorsProtocolPrivate):
return PreprocessingResponse(status=True, results=chunks)
def _resolve_chunk_size_params(self, options: PreprocessorOptions) -> Tuple[int, int]:
window_len = options.get(str(SimpleChunkingOptions.chunk_size_in_tokens), self.config.chunk_size_in_tokens)
window_len = (options or {}).get(
str(SimpleChunkingOptions.chunk_size_in_tokens), self.config.chunk_size_in_tokens
)
chunk_overlap_ratio = options.get(
chunk_overlap_ratio = (options or {}).get(
str(SimpleChunkingOptions.chunk_overlap_ratio), self.config.chunk_overlap_ratio
)
overlap_len = window_len // chunk_overlap_ratio