AVB/bedrock-access-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added missing input type checking (Cohere Embed); Removed fingerprint from BaseEmbeddingsResponse

Browse Source
This commit is contained in:
Joao Galego
2024-03-28 08:05:39 +00:00
parent b24a43f6f4
commit c3b7395028
2 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/api/models/bedrock.py
View File

@@ -483,8 +483,12 @@ def get_embeddings_model(model_id: str) -> BedrockEmbeddingsModel:
class CohereEmbeddingsModel(BedrockEmbeddingsModel): class CohereEmbeddingsModel(BedrockEmbeddingsModel):
def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict: def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict:
if isinstance(embeddings_request.input, str):
texts = [embeddings_request.input]
elif isinstance(embeddings_request.input, list):
texts = embeddings_request.input
args = { args = {
"texts": embeddings_request.input, "texts": texts,
"input_type": embeddings_request.input_type if embeddings_request.input_type else "search_document", "input_type": embeddings_request.input_type if embeddings_request.input_type else "search_document",
"truncate": embeddings_request.truncate if embeddings_request.truncate else "NONE", "truncate": embeddings_request.truncate if embeddings_request.truncate else "NONE",
} }

1
src/api/schema.py
View File

@@ -93,7 +93,6 @@ class EmbeddingsRequest(BaseModel):
class BaseEmbeddingsResponse(BaseModel): class BaseEmbeddingsResponse(BaseModel):
created: int = Field(default_factory=lambda: int(time.time())) created: int = Field(default_factory=lambda: int(time.time()))
model: str model: str
system_fingerprint: str = "fp_e97c09dd4e26"
class Embedding(BaseModel): class Embedding(BaseModel):