mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
Dynamic embeddings model (#556)
* Dynamic embeddings model selection * Added tests * HF embeddings are skipped, tests don't run with that package currently tests
This commit is contained in:
parent
6129bb68c1
commit
d9d4c91363
8 changed files with 816 additions and 11 deletions
|
|
@ -9,7 +9,7 @@ from prometheus_client import Histogram
|
|||
|
||||
from .. schema import EmbeddingsRequest, EmbeddingsResponse, Error
|
||||
from .. exceptions import TooManyRequests
|
||||
from .. base import FlowProcessor, ConsumerSpec, ProducerSpec
|
||||
from .. base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec
|
||||
|
||||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -45,6 +45,12 @@ class EmbeddingsService(FlowProcessor):
|
|||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ParameterSpec(
|
||||
name = "model",
|
||||
)
|
||||
)
|
||||
|
||||
async def on_request(self, msg, consumer, flow):
|
||||
|
||||
try:
|
||||
|
|
@ -57,7 +63,9 @@ class EmbeddingsService(FlowProcessor):
|
|||
|
||||
logger.debug(f"Handling embeddings request {id}...")
|
||||
|
||||
vectors = await self.on_embeddings(request.text)
|
||||
# Pass model from request if specified (non-empty), otherwise use default
|
||||
model = flow("model")
|
||||
vectors = await self.on_embeddings(request.text, model=model)
|
||||
|
||||
await flow("response").send(
|
||||
EmbeddingsResponse(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue