mseep-txtai 9.1.1__py3-none-any.whl

txtai/api/routers/openai.py

@@ -0,0 +1,191 @@
+"""
+Defines an OpenAI-compatible API endpoint for txtai.
+
+See the following specification for more information:
+https://github.com/openai/openai-openapi
+"""
+
+import uuid
+import json
+import time
+
+from typing import List, Optional, Union
+
+from fastapi import APIRouter, Body, Form, UploadFile
+from fastapi.responses import Response, StreamingResponse
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+# pylint: disable=W0622
+@router.post("/v1/chat/completions")
+def chat(
+    messages: List[dict] = Body(...),
+    model: str = Body(...),
+    max_completion_tokens: Optional[int] = Body(default=None),
+    stream: Optional[bool] = Body(default=False),
+):
+    """
+    Runs a chat completion request.
+
+    Args:
+        messages: list of messages [{"role": role, "content": content}]
+        model: agent name, workflow name, pipeline name or embeddings
+        max_completion_tokens: sets the max length to generate
+        stream: streams response if True
+
+    Returns:
+        chat completion
+    """
+
+    # Build keyword arguments
+    kwargs = {key: value for key, value in [("stream", stream), ("maxlength", max_completion_tokens)] if value}
+
+    # Get first message
+    message = messages[0]["content"]
+
+    # Agent
+    if model in application.get().agents:
+        result = application.get().agent(model, message, **kwargs)
+
+    # Embeddings search
+    elif model == "embeddings":
+        result = application.get().search(message, 1, **kwargs)[0]["text"]
+
+    # Pipeline
+    elif model in application.get().pipelines and model != "llm":
+        result = application.get().pipeline(model, message, **kwargs)
+
+    # Workflow
+    elif model in application.get().workflows:
+        result = list(application.get().workflow(model, [message], **kwargs))[0]
+
+    # Default to running all messages through default LLM
+    else:
+        result = application.get().pipeline("llm", messages, **kwargs)
+
+    # Write response
+    return StreamingResponse(StreamingChatResponse()(model, result)) if stream else ChatResponse()(model, result)
+
+
+@router.post("/v1/embeddings")
+def embeddings(input: Union[str, List[str]] = Body(...), model: str = Body(...)):
+    """
+    Creates an embeddings vector for the input text.
+
+    Args:
+        input: text|list
+        model: model name
+
+    Returns:
+        list of embeddings vectors
+    """
+
+    # Convert to embeddings
+    result = application.get().batchtransform([input] if isinstance(input, str) else input)
+
+    # Build and return response
+    data = []
+    for index, embedding in enumerate(result):
+        data.append({"object": "embedding", "embedding": embedding, "index": index})
+
+    return {"object": "list", "data": data, "model": model}
+
+
+@router.post("/v1/audio/speech")
+def speech(input: str = Body(...), voice: str = Body(...), response_format: Optional[str] = Body(default="mp3")):
+    """
+    Generates speech for the input text.
+
+    Args:
+        input: input text
+        voice: speaker name
+        response_format: audio encoding format, defaults to mp3
+
+    Returns:
+        audio data
+    """
+
+    # Convert to audio
+    audio = application.get().pipeline("texttospeech", input, speaker=voice, encoding=response_format)
+
+    # Write audio
+    return Response(audio)
+
+
+@router.post("/v1/audio/transcriptions")
+def transcribe(file: UploadFile, language: Optional[str] = Form(default=None), response_format: Optional[str] = Form(default="json")):
+    """
+    Transcribes audio to text.
+
+    Args:
+        file: audio input file
+        language: language of input audio
+        response_format: output format (json or text)
+
+    Returns:
+        transcribed text
+    """
+
+    # Transcribe
+    text = application.get().pipeline("transcription", file.file, language=language, task="transcribe")
+    return text if response_format == "text" else {"text": text}
+
+
+@router.post("/v1/audio/translations")
+def translate(
+    file: UploadFile,
+    response_format: Optional[str] = Form(default="json"),
+):
+    """
+    Translates audio to English.
+
+    Args:
+        file: audio input file
+        response_format: output format (json or text)
+
+    Returns:
+        translated text
+    """
+
+    # Transcribe and translate to English
+    text = application.get().pipeline("transcription", file.file, language="English", task="translate")
+    return text if response_format == "text" else {"text": text}
+
+
+class ChatResponse:
+    """
+    Returns a chat response object.
+    """
+
+    def __call__(self, model, result):
+        return {
+            "id": str(uuid.uuid4()),
+            "object": "chat.completion",
+            "created": int(time.time() * 1000),
+            "model": model,
+            "choices": [{"id": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
+        }
+
+
+class StreamingChatResponse:
+    """
+    Returns a streaming chat response object.
+    """
+
+    def __call__(self, model, result):
+        for chunk in result:
+            yield "data: " + json.dumps(
+                {
+                    "id": str(uuid.uuid4()),
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time() * 1000),
+                    "model": model,
+                    "choices": [{"id": 0, "delta": {"content": chunk}}],
+                }
+            ) + "\n\n"
+
+        yield "data: [DONE]\n\n"

txtai/api/routers/rag.py

@@ -0,0 +1,61 @@
+"""
+Defines API paths for rag endpoints.
+"""
+
+from typing import List, Optional
+
+from fastapi import APIRouter, Body
+from fastapi.responses import StreamingResponse
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/rag")
+def rag(query: str, maxlength: Optional[int] = None, stream: Optional[bool] = False):
+    """
+    Runs a RAG pipeline for the input query.
+
+    Args:
+        query: input RAG query
+        maxlength: optional response max length
+        stream: streams response if True
+
+    Returns:
+        answer
+    """
+
+    # Build keyword arguments
+    kwargs = {key: value for key, value in [("stream", stream), ("maxlength", maxlength)] if value}
+
+    # Run pipeline
+    result = application.get().pipeline("rag", query, **kwargs)
+
+    # Handle both standard and streaming responses
+    return StreamingResponse(result) if stream else result
+
+
+@router.post("/batchrag")
+def batchrag(queries: List[str] = Body(...), maxlength: Optional[int] = Body(default=None), stream: Optional[bool] = Body(default=False)):
+    """
+    Runs a RAG pipeline for the input queries.
+
+    Args:
+        queries: input RAG queries
+        maxlength: optional response max length
+        stream: streams response if True
+
+    Returns:
+        answers
+    """
+
+    # Build keyword arguments
+    kwargs = {key: value for key, value in [("stream", stream), ("maxlength", maxlength)] if value}
+
+    # Run pipeline
+    result = application.get().pipeline("rag", queries, **kwargs)
+
+    # Handle both standard and streaming responses
+    return StreamingResponse(result) if stream else result

txtai/api/routers/reranker.py

@@ -0,0 +1,46 @@
+"""
+Defines API paths for reranking endpoints.
+"""
+
+from typing import List, Optional
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/rerank")
+def rerank(query: str, limit: Optional[int] = 3, factor: Optional[int] = 10):
+    """
+    Queries an embeddings database and reranks the results with a similarity pipeline.
+
+    Args:
+        query: query text
+        limit: maximum results
+        factor: factor to multiply limit by for the initial embeddings search
+
+    Returns:
+        query results
+    """
+
+    return application.get().pipeline("reranker", (query, limit, factor))
+
+
+@router.post("/batchrerank")
+def batchrerank(queries: List[str] = Body(...), limit: Optional[int] = Body(default=3), factor: Optional[int] = Body(default=10)):
+    """
+    Queries an embeddings database and reranks the results with a similarity pipeline.
+
+    Args:
+        queries: list of queries
+        limit: maximum results
+        factor: factor to multiply limit by for the initial embeddings search
+
+    Returns:
+        query results
+    """
+
+    return application.get().pipeline("reranker", (queries, limit, factor))

txtai/api/routers/segmentation.py

@@ -0,0 +1,42 @@
+"""
+Defines API paths for segmentation endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/segment")
+def segment(text: str):
+    """
+    Segments text into semantic units.
+
+    Args:
+        text: input text
+
+    Returns:
+        segmented text
+    """
+
+    return application.get().pipeline("segmentation", (text,))
+
+
+@router.post("/batchsegment")
+def batchsegment(texts: List[str] = Body(...)):
+    """
+    Segments text into semantic units.
+
+    Args:
+        texts: list of texts to segment
+
+    Returns:
+        list of segmented text
+    """
+
+    return application.get().pipeline("segmentation", (texts,))

txtai/api/routers/similarity.py

@@ -0,0 +1,48 @@
+"""
+Defines API paths for similarity endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.post("/similarity")
+def similarity(query: str = Body(...), texts: List[str] = Body(...)):
+    """
+    Computes the similarity between query and list of text. Returns a list of
+    {id: value, score: value} sorted by highest score, where id is the index
+    in texts.
+
+    Args:
+        query: query text
+        texts: list of text
+
+    Returns:
+        list of {id: value, score: value}
+    """
+
+    return application.get().similarity(query, texts)
+
+
+@router.post("/batchsimilarity")
+def batchsimilarity(queries: List[str] = Body(...), texts: List[str] = Body(...)):
+    """
+    Computes the similarity between list of queries and list of text. Returns a list
+    of {id: value, score: value} sorted by highest score per query, where id is the
+    index in texts.
+
+    Args:
+        queries: queries text
+        texts: list of text
+
+    Returns:
+        list of {id: value, score: value} per query
+    """
+
+    return application.get().batchsimilarity(queries, texts)

txtai/api/routers/summary.py

@@ -0,0 +1,46 @@
+"""
+Defines API paths for summary endpoints.
+"""
+
+from typing import List, Optional
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/summary")
+def summary(text: str, minlength: Optional[int] = None, maxlength: Optional[int] = None):
+    """
+    Runs a summarization model against a block of text.
+
+    Args:
+        text: text to summarize
+        minlength: minimum length for summary
+        maxlength: maximum length for summary
+
+    Returns:
+        summary text
+    """
+
+    return application.get().pipeline("summary", (text, minlength, maxlength))
+
+
+@router.post("/batchsummary")
+def batchsummary(texts: List[str] = Body(...), minlength: Optional[int] = Body(default=None), maxlength: Optional[int] = Body(default=None)):
+    """
+    Runs a summarization model against a block of text.
+
+    Args:
+        texts: list of text to summarize
+        minlength: minimum length for summary
+        maxlength: maximum length for summary
+
+    Returns:
+        list of summary text
+    """
+
+    return application.get().pipeline("summary", (texts, minlength, maxlength))

txtai/api/routers/tabular.py

@@ -0,0 +1,42 @@
+"""
+Defines API paths for tabular endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/tabular")
+def tabular(file: str):
+    """
+    Splits tabular data into rows and columns.
+
+    Args:
+        file: file to process
+
+    Returns:
+        list of (id, text, tag) elements
+    """
+
+    return application.get().pipeline("tabular", (file,))
+
+
+@router.post("/batchtabular")
+def batchtabular(files: List[str] = Body(...)):
+    """
+    Splits tabular data into rows and columns.
+
+    Args:
+        files: list of files to process
+
+    Returns:
+        list of (id, text, tag) elements
+    """
+
+    return application.get().pipeline("tabular", (files,))

txtai/api/routers/textractor.py

@@ -0,0 +1,42 @@
+"""
+Defines API paths for textractor endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/textract")
+def textract(file: str):
+    """
+    Extracts text from a file at path.
+
+    Args:
+        file: file to extract text
+
+    Returns:
+        extracted text
+    """
+
+    return application.get().pipeline("textractor", (file,))
+
+
+@router.post("/batchtextract")
+def batchtextract(files: List[str] = Body(...)):
+    """
+    Extracts text from a file at path.
+
+    Args:
+        files: list of files to extract text
+
+    Returns:
+        list of extracted text
+    """
+
+    return application.get().pipeline("textractor", (files,))

txtai/api/routers/texttospeech.py

@@ -0,0 +1,33 @@
+"""
+Defines API paths for TTS endpoints
+"""
+
+from typing import Optional
+
+from fastapi import APIRouter, Response
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/texttospeech")
+def texttospeech(text: str, speaker: Optional[str] = None, encoding: Optional[str] = "mp3"):
+    """
+    Generates speech from text.
+
+    Args:
+        text: text
+        speaker: speaker id, defaults to 1
+        encoding: optional audio encoding format
+
+    Returns:
+        Audio data
+    """
+
+    # Convert to audio
+    audio = application.get().pipeline("texttospeech", text, speaker=speaker, encoding=encoding)
+
+    # Write audio
+    return Response(audio, headers={"Content-Disposition": f"attachment;filename=speech.{encoding.lower()}"})

txtai/api/routers/transcription.py

@@ -0,0 +1,42 @@
+"""
+Defines API paths for transcription endpoints.
+"""
+
+from typing import List
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/transcribe")
+def transcribe(file: str):
+    """
+    Transcribes audio files to text.
+
+    Args:
+        file: file to transcribe
+
+    Returns:
+        transcribed text
+    """
+
+    return application.get().pipeline("transcription", (file,))
+
+
+@router.post("/batchtranscribe")
+def batchtranscribe(files: List[str] = Body(...)):
+    """
+    Transcribes audio files to text.
+
+    Args:
+        files: list of files to transcribe
+
+    Returns:
+        list of transcribed text
+    """
+
+    return application.get().pipeline("transcription", (files,))

txtai/api/routers/translation.py

@@ -0,0 +1,46 @@
+"""
+Defines API paths for translation endpoints.
+"""
+
+from typing import List, Optional
+
+from fastapi import APIRouter, Body
+
+from .. import application
+from ..route import EncodingAPIRoute
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.get("/translate")
+def translate(text: str, target: Optional[str] = "en", source: Optional[str] = None):
+    """
+    Translates text from source language into target language.
+
+    Args:
+        text: text to translate
+        target: target language code, defaults to "en"
+        source: source language code, detects language if not provided
+
+    Returns:
+        translated text
+    """
+
+    return application.get().pipeline("translation", (text, target, source))
+
+
+@router.post("/batchtranslate")
+def batchtranslate(texts: List[str] = Body(...), target: Optional[str] = Body(default="en"), source: Optional[str] = Body(default=None)):
+    """
+    Translates text from source language into target language.
+
+    Args:
+        texts: list of text to translate
+        target: target language code, defaults to "en"
+        source: source language code, detects language if not provided
+
+    Returns:
+        list of translated text
+    """
+
+    return application.get().pipeline("translation", (texts, target, source))

txtai/api/routers/upload.py

@@ -0,0 +1,36 @@
+"""
+Defines API paths for upload endpoints.
+"""
+
+import shutil
+import tempfile
+
+from typing import List
+
+from fastapi import APIRouter, File, Form, UploadFile
+
+from ..route import EncodingAPIRoute
+
+
+router = APIRouter(route_class=EncodingAPIRoute)
+
+
+@router.post("/upload")
+def upload(files: List[UploadFile] = File(), suffix: str = Form(default=None)):
+    """
+    Uploads files for local server processing.
+
+    Args:
+        data: list of files to upload
+
+    Returns:
+        list of server paths
+    """
+
+    paths = []
+    for f in files:
+        with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=suffix) as tmp:
+            shutil.copyfileobj(f.file, tmp)
+            paths.append(tmp.name)
+
+    return paths