llmflows

llmflows copied to clipboard

https://github.com/stoyan-stoyanov/llmflows/blob/6ce79a8cc3b06e799d2b6eb6f6e8e78354f68dcd/llmflows/llms/openai_embeddings.py#L50-L52

I am not a fan of functions that can return either one or the other type, however, I assume that given a List[VectorDoc], generate will always return a List[VectorDoc]. If so, then you can use the @overload decorator and declare two function signatures:

from typing import overload

@overload
def generate(docs: VectorDoc) -> VectorDoc:
    ...

@overload
def generate(docs: List[VectorDoc]) -> List[VectorDoc]:
    ...

def generate(x):
    # here goes the implementation

Note that the ellipsis ... is required.