diff --git a/semantic_router/index/base.py b/semantic_router/index/base.py index d25d41dc663487959cf6da5b5ce8e365010b1a0a..3d3910834f0912307ab69d1b21e5d1ef5bb38d6f 100644 --- a/semantic_router/index/base.py +++ b/semantic_router/index/base.py @@ -26,7 +26,7 @@ class BaseIndex(BaseModel): embeddings: List[List[float]], routes: List[str], utterances: List[Any], - function_schemas: List[Dict[str, Any]], + function_schemas: List[Dict[str, Any]] = None, # type: ignore ): """ Add embeddings to the index. diff --git a/semantic_router/index/local.py b/semantic_router/index/local.py index 7150b267587715d09124eb6f94ab242a7795c3f9..802455cb77fbf1d18ac07929508b190684d984f4 100644 --- a/semantic_router/index/local.py +++ b/semantic_router/index/local.py @@ -5,6 +5,7 @@ import numpy as np from semantic_router.index.base import BaseIndex from semantic_router.linear import similarity_matrix, top_scores from semantic_router.utils.logger import logger +from typing import Any class LocalIndex(BaseIndex): @@ -26,6 +27,7 @@ class LocalIndex(BaseIndex): embeddings: List[List[float]], routes: List[str], utterances: List[str], + function_schemas: List[Dict[str, Any]] = None, # type: ignore ): embeds = np.array(embeddings) # type: ignore routes_arr = np.array(routes) diff --git a/semantic_router/index/pinecone.py b/semantic_router/index/pinecone.py index 0b84ad4cd3a061b11c09b3c7335f75ff13ad7a0c..da11ec6ea4df8ee762fb2d84b1a66a0de504d0c2 100644 --- a/semantic_router/index/pinecone.py +++ b/semantic_router/index/pinecone.py @@ -310,7 +310,7 @@ class PineconeIndex(BaseIndex): embeddings: List[List[float]], routes: List[str], utterances: List[str], - function_schemas: List[Dict[str, Any]] = "", + function_schemas: List[Dict[str, Any]] = None, # type: ignore batch_size: int = 100, ): """Add vectors to Pinecone in batches.""" diff --git a/semantic_router/index/postgres.py b/semantic_router/index/postgres.py index 4c971d4d1a71ea51a05f0698b4f7aa971aeeef5d..9fbac62fde49fe53905e27e2d78e4eb8c67af8c6 100644 --- a/semantic_router/index/postgres.py +++ b/semantic_router/index/postgres.py @@ -254,7 +254,11 @@ class PostgresIndex(BaseIndex): raise ValueError("No comment found for the 'vector' column.") def add( - self, embeddings: List[List[float]], routes: List[str], utterances: List[Any] + self, + embeddings: List[List[float]], + routes: List[str], + utterances: List[Any], + function_schemas: List[Dict[str, Any]] = None, # type: ignore ) -> None: """ Adds vectors to the index. diff --git a/semantic_router/index/qdrant.py b/semantic_router/index/qdrant.py index c1a5e28b58792091e4fb568248827954547e37a5..0fc6aa52c1ee8be1d769a8463ff68a6e31e3bfe5 100644 --- a/semantic_router/index/qdrant.py +++ b/semantic_router/index/qdrant.py @@ -175,6 +175,7 @@ class QdrantIndex(BaseIndex): embeddings: List[List[float]], routes: List[str], utterances: List[str], + function_schemas: List[Dict[str, Any]] = None, # type: ignore batch_size: int = DEFAULT_UPLOAD_BATCH_SIZE, ): self.dimensions = self.dimensions or len(embeddings[0]) diff --git a/semantic_router/layer.py b/semantic_router/layer.py index 2285a7d9c00a1e87c042069613e83e24fc5a873c..bae67544940c9a30333ae7331a0abcf354873ab8 100644 --- a/semantic_router/layer.py +++ b/semantic_router/layer.py @@ -437,7 +437,7 @@ class RouteLayer: function_schemas=( route.function_schemas * len(route.utterances) if route.function_schemas - else [""] * len(route.utterances) + else [""] * len(route.utterances) # type: ignore ), ) @@ -482,7 +482,9 @@ class RouteLayer: def _add_routes(self, routes: List[Route]): # create embeddings for all routes - route_names, all_utterances = self._extract_routes_details(routes) + route_names, all_utterances, function_schemas = self._extract_routes_details( + routes + ) embedded_utterances = self.encoder(all_utterances) # create route array # add everything to the index @@ -490,11 +492,14 @@ class RouteLayer: embeddings=embedded_utterances, routes=route_names, utterances=all_utterances, + function_schemas=function_schemas, ) def _add_and_sync_routes(self, routes: List[Route]): # create embeddings for all routes and sync at startup with remote ones based on sync setting - local_route_names, local_utterances = self._extract_routes_details(routes) + local_route_names, local_utterances, local_function_schemas = ( + self._extract_routes_details(routes) + ) routes_to_add, routes_to_delete, layer_routes_dict = self.index._sync_index( local_route_names=local_route_names, local_utterances=local_utterances, @@ -522,6 +527,7 @@ class RouteLayer: embeddings=embedded_utterances_to_add, routes=route_names_to_add, utterances=all_utterances_to_add, + function_schemas=local_function_schemas, ) self._set_layer_routes(layer_routes) @@ -529,7 +535,13 @@ class RouteLayer: def _extract_routes_details(self, routes: List[Route]) -> Tuple: route_names = [route.name for route in routes for _ in route.utterances] utterances = [utterance for route in routes for utterance in route.utterances] - return route_names, utterances + function_schemas = [ + function_schema if function_schema is not None else "" + for route in routes + if route.function_schemas is not None + for function_schema in route.function_schemas + ] + return route_names, utterances, function_schemas def _encode(self, text: str) -> Any: """Given some text, encode it."""