Skip to content
Snippets Groups Projects
Commit f3d2bcd2 authored by Simonas's avatar Simonas
Browse files

merge main

parents d10c2797 da1596a2
No related branches found
No related tags found
No related merge requests found
<?xml version="1.0" ?> <?xml version="1.0" ?>
<coverage version="7.3.2" timestamp="1702633916069" lines-valid="344" lines-covered="344" line-rate="1" branches-covered="0" branches-valid="0" branch-rate="0" complexity="0"> <coverage version="7.3.3" timestamp="1702893702032" lines-valid="345" lines-covered="345" line-rate="1" branches-covered="0" branches-valid="0" branch-rate="0" complexity="0">
<!-- Generated by coverage.py: https://coverage.readthedocs.io/en/7.3.2 --> <!-- Generated by coverage.py: https://coverage.readthedocs.io/en/7.3.3 -->
<!-- Based on https://raw.githubusercontent.com/cobertura/web/master/htdocs/xml/coverage-04.dtd --> <!-- Based on https://raw.githubusercontent.com/cobertura/web/master/htdocs/xml/coverage-04.dtd -->
<sources> <sources>
<source>/Users/jakit/customers/aurelio/semantic-router/semantic_router</source> <source>/Users/jakit/customers/aurelio/semantic-router/semantic_router</source>
...@@ -382,8 +382,9 @@ ...@@ -382,8 +382,9 @@
<line number="47" hits="1"/> <line number="47" hits="1"/>
<line number="49" hits="1"/> <line number="49" hits="1"/>
<line number="50" hits="1"/> <line number="50" hits="1"/>
<line number="52" hits="1"/> <line number="51" hits="1"/>
<line number="53" hits="1"/> <line number="53" hits="1"/>
<line number="54" hits="1"/>
</lines> </lines>
</class> </class>
</classes> </classes>
......
File added
File added
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
## Define LLMs ## Define LLMs
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
%reload_ext dotenv
%dotenv
```
%% Cell type:code id: tags:
``` python
# OpenAI # OpenAI
import os import os
import openai import openai
from semantic_router.utils.logger import logger from semantic_router.utils.logger import logger
# Docs # https://platform.openai.com/docs/guides/function-calling # Docs # https://platform.openai.com/docs/guides/function-calling
def llm_openai(prompt: str, model: str = "gpt-4") -> str: def llm_openai(prompt: str, model: str = "gpt-4") -> str:
try: try:
logger.info(f"Calling {model} model") logger.info(f"Calling {model} model")
response = openai.chat.completions.create( response = openai.chat.completions.create(
model=model, model=model,
messages=[ messages=[
{"role": "system", "content": f"{prompt}"}, {"role": "system", "content": f"{prompt}"},
], ],
) )
ai_message = response.choices[0].message.content ai_message = response.choices[0].message.content
if not ai_message: if not ai_message:
raise Exception("AI message is empty", ai_message) raise Exception("AI message is empty", ai_message)
logger.info(f"AI message: {ai_message}") logger.info(f"AI message: {ai_message}")
return ai_message return ai_message
except Exception as e: except Exception as e:
raise Exception("Failed to call OpenAI API", e) raise Exception("Failed to call OpenAI API", e)
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
# Mistral # Mistral
import os import os
import requests import requests
# Docs https://huggingface.co/docs/transformers/main_classes/text_generation # Docs https://huggingface.co/docs/transformers/main_classes/text_generation
HF_API_TOKEN = os.getenv("HF_API_TOKEN") HF_API_TOKEN = os.getenv("HF_API_TOKEN")
def llm_mistral(prompt: str) -> str: def llm_mistral(prompt: str) -> str:
api_url = "https://z5t4cuhg21uxfmc3.us-east-1.aws.endpoints.huggingface.cloud/" api_url = "https://z5t4cuhg21uxfmc3.us-east-1.aws.endpoints.huggingface.cloud/"
headers = { headers = {
"Authorization": f"Bearer {HF_API_TOKEN}", "Authorization": f"Bearer {HF_API_TOKEN}",
"Content-Type": "application/json", "Content-Type": "application/json",
} }
logger.info("Calling Mistral model") logger.info("Calling Mistral model")
response = requests.post( response = requests.post(
api_url, api_url,
headers=headers, headers=headers,
json={ json={
"inputs": f"You are a helpful assistant, user query: {prompt}", "inputs": f"You are a helpful assistant, user query: {prompt}",
"parameters": { "parameters": {
"max_new_tokens": 200, "max_new_tokens": 200,
"temperature": 0.1, "temperature": 0.1,
}, },
}, },
) )
if response.status_code != 200: if response.status_code != 200:
raise Exception("Failed to call HuggingFace API", response.text) raise Exception("Failed to call HuggingFace API", response.text)
ai_message = response.json()[0]["generated_text"] ai_message = response.json()[0]["generated_text"]
if not ai_message: if not ai_message:
raise Exception("AI message is empty", ai_message) raise Exception("AI message is empty", ai_message)
logger.info(f"AI message: {ai_message}") logger.info(f"AI message: {ai_message}")
return ai_message return ai_message
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Now we need to generate config from function schema using LLM ### Now we need to generate config from function schema using LLM
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
import inspect import inspect
from typing import Any from typing import Any
def get_function_schema(function) -> dict[str, Any]: def get_function_schema(function) -> dict[str, Any]:
schema = { schema = {
"name": function.__name__, "name": function.__name__,
"description": str(inspect.getdoc(function)), "description": str(inspect.getdoc(function)),
"signature": str(inspect.signature(function)), "signature": str(inspect.signature(function)),
"output": str( "output": str(
inspect.signature(function).return_annotation, inspect.signature(function).return_annotation,
), ),
} }
return schema return schema
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
import json import json
from semantic_router.utils.logger import logger from semantic_router.utils.logger import logger
def generate_route(function) -> dict: def generate_route(function) -> dict:
logger.info("Generating config...") logger.info("Generating config...")
example_schema = { example_schema = {
"name": "get_weather", "name": "get_weather",
"description": "Useful to get the weather in a specific location", "description": "Useful to get the weather in a specific location",
"signature": "(location: str) -> str", "signature": "(location: str) -> str",
"output": "<class 'str'>", "output": "<class 'str'>",
} }
example_config = { example_config = {
"name": "get_weather", "name": "get_weather",
"utterances": [ "utterances": [
"What is the weather like in SF?", "What is the weather like in SF?",
"What is the weather in Cyprus?", "What is the weather in Cyprus?",
"weather in London?", "weather in London?",
"Tell me the weather in New York", "Tell me the weather in New York",
"what is the current weather in Paris?", "what is the current weather in Paris?",
], ],
} }
function_schema = get_function_schema(function) function_schema = get_function_schema(function)
prompt = f""" prompt = f"""
You are a helpful assistant designed to output JSON. You are a helpful assistant designed to output JSON.
Given the following function schema Given the following function schema
{function_schema} {function_schema}
generate a routing config with the format: generate a routing config with the format:
{example_config} {example_config}
For example: For example:
Input: {example_schema} Input: {example_schema}
Output: {example_config} Output: {example_config}
Input: {function_schema} Input: {function_schema}
Output: Output:
""" """
ai_message = llm_openai(prompt) ai_message = llm_openai(prompt)
ai_message = ai_message.replace("CONFIG:", "").replace("'", '"').strip().rstrip(",") ai_message = ai_message.replace("CONFIG:", "").replace("'", '"').strip().rstrip(",")
try: try:
route_config = json.loads(ai_message) route_config = json.loads(ai_message)
logger.info(f"Generated config: {route_config}") logger.info(f"Generated config: {route_config}")
return route_config return route_config
except json.JSONDecodeError as json_error: except json.JSONDecodeError as json_error:
logger.error(f"JSON parsing error {json_error}") logger.error(f"JSON parsing error {json_error}")
print(f"AI message: {ai_message}") print(f"AI message: {ai_message}")
return {"error": "Failed to generate config"} return {"error": "Failed to generate config"}
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Extract function parameters using `Mistral` open-source model Extract function parameters using `Mistral` open-source model
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
def extract_parameters(query: str, function) -> dict: def extract_parameters(query: str, function) -> dict:
logger.info("Extracting parameters...") logger.info("Extracting parameters...")
example_query = "How is the weather in Hawaii right now in International units?" example_query = "How is the weather in Hawaii right now in International units?"
example_schema = { example_schema = {
"name": "get_weather", "name": "get_weather",
"description": "Useful to get the weather in a specific location", "description": "Useful to get the weather in a specific location",
"signature": "(location: str, degree: str) -> str", "signature": "(location: str, degree: str) -> str",
"output": "<class 'str'>", "output": "<class 'str'>",
} }
example_parameters = { example_parameters = {
"location": "London", "location": "London",
"degree": "Celsius", "degree": "Celsius",
} }
prompt = f""" prompt = f"""
You are a helpful assistant designed to output JSON. You are a helpful assistant designed to output JSON.
Given the following function schema Given the following function schema
{get_function_schema(function)} {get_function_schema(function)}
and query and query
{query} {query}
extract the parameters values from the query, in a valid JSON format. extract the parameters values from the query, in a valid JSON format.
Example: Example:
Input: Input:
query: {example_query} query: {example_query}
schema: {example_schema} schema: {example_schema}
Output: Output:
parameters: {example_parameters} parameters: {example_parameters}
Input: Input:
query: {query} query: {query}
schema: {get_function_schema(function)} schema: {get_function_schema(function)}
Output: Output:
parameters: parameters:
""" """
ai_message = llm_mistral(prompt) ai_message = llm_mistral(prompt)
ai_message = ai_message.replace("CONFIG:", "").replace("'", '"').strip().rstrip(",") ai_message = ai_message.replace("CONFIG:", "").replace("'", '"').strip().rstrip(",")
try: try:
parameters = json.loads(ai_message) parameters = json.loads(ai_message)
logger.info(f"Extracted parameters: {parameters}") logger.info(f"Extracted parameters: {parameters}")
return parameters return parameters
except json.JSONDecodeError as json_error: except json.JSONDecodeError as json_error:
logger.error(f"JSON parsing error {json_error}") logger.error(f"JSON parsing error {json_error}")
return {"error": "Failed to extract parameters"} return {"error": "Failed to extract parameters"}
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Set up the routing layer Set up the routing layer
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
from semantic_router.schema import Route from semantic_router.schema import Route
from semantic_router.encoders import CohereEncoder, OpenAIEncoder from semantic_router.encoders import CohereEncoder, OpenAIEncoder
from semantic_router.layer import RouteLayer from semantic_router.layer import RouteLayer
from semantic_router.utils.logger import logger from semantic_router.utils.logger import logger
def create_router(routes: list[dict]) -> RouteLayer: def create_router(routes: list[dict]) -> RouteLayer:
logger.info("Creating route layer...") logger.info("Creating route layer...")
encoder = OpenAIEncoder encoder = OpenAIEncoder
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
from semantic_router.schema import Route from semantic_router.schema import Route
from semantic_router.encoders import CohereEncoder from semantic_router.encoders import CohereEncoder
from semantic_router.layer import RouteLayer from semantic_router.layer import RouteLayer
from semantic_router.utils.logger import logger from semantic_router.utils.logger import logger
def create_router(routes: list[dict]) -> RouteLayer: def create_router(routes: list[dict]) -> RouteLayer:
logger.info("Creating route layer...") logger.info("Creating route layer...")
encoder = OpenAIEncoder() encoder = OpenAIEncoder()
route_list: list[Route] = [] route_list: list[Route] = []
for route in routes: for route in routes:
if "name" in route and "utterances" in route: if "name" in route and "utterances" in route:
print(f"Route: {route}") print(f"Route: {route}")
route_list.append(Route(name=route["name"], utterances=route["utterances"])) route_list.append(Route(name=route["name"], utterances=route["utterances"]))
else: else:
logger.warning(f"Misconfigured route: {route}") logger.warning(f"Misconfigured route: {route}")
return RouteLayer(encoder=encoder, routes=route_list) return RouteLayer(encoder=encoder, routes=route_list)
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
Set up calling functions Set up calling functions
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
from typing import Callable from typing import Callable
def call_function(function: Callable, parameters: dict[str, str]): def call_function(function: Callable, parameters: dict[str, str]):
try: try:
return function(**parameters) return function(**parameters)
except TypeError as e: except TypeError as e:
logger.error(f"Error calling function: {e}") logger.error(f"Error calling function: {e}")
def call_llm(query: str): def call_llm(query: str):
return llm_mistral(query) return llm_mistral(query)
def call(query: str, functions: list[Callable], router: RouteLayer): def call(query: str, functions: list[Callable], router: RouteLayer):
function_name = router(query) function_name = router(query)
if not function_name: if not function_name:
logger.warning("No function found") logger.warning("No function found")
return call_llm(query) return call_llm(query)
for function in functions: for function in functions:
if function.__name__ == function_name: if function.__name__ == function_name:
parameters = extract_parameters(query, function) parameters = extract_parameters(query, function)
print(f"parameters: {parameters}") print(f"parameters: {parameters}")
return call_function(function, parameters) return call_function(function, parameters)
``` ```
%% Cell type:markdown id: tags: %% Cell type:markdown id: tags:
### Workflow ### Workflow
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
def get_time(location: str) -> str: def get_time(location: str) -> str:
"""Useful to get the time in a specific location""" """Useful to get the time in a specific location"""
print(f"Calling `get_time` function with location: {location}") print(f"Calling `get_time` function with location: {location}")
return "get_time" return "get_time"
def get_news(category: str, country: str) -> str: def get_news(category: str, country: str) -> str:
"""Useful to get the news in a specific country""" """Useful to get the news in a specific country"""
print( print(
f"Calling `get_news` function with category: {category} and country: {country}" f"Calling `get_news` function with category: {category} and country: {country}"
) )
return "get_news" return "get_news"
# Registering functions to the router # Registering functions to the router
route_get_time = generate_route(get_time) route_get_time = generate_route(get_time)
route_get_news = generate_route(get_news) route_get_news = generate_route(get_news)
routes = [route_get_time, route_get_news] routes = [route_get_time, route_get_news]
router = create_router(routes) router = create_router(routes)
# Tools # Tools
tools = [get_time, get_news] tools = [get_time, get_news]
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
def get_time(location: str) -> str: def get_time(location: str) -> str:
"""Useful to get the time in a specific location""" """Useful to get the time in a specific location"""
print(f"Calling `get_time` function with location: {location}") print(f"Calling `get_time` function with location: {location}")
return "get_time" return "get_time"
def get_news(category: str, country: str) -> str: def get_news(category: str, country: str) -> str:
"""Useful to get the news in a specific country""" """Useful to get the news in a specific country"""
print( print(
f"Calling `get_news` function with category: {category} and country: {country}" f"Calling `get_news` function with category: {category} and country: {country}"
) )
return "get_news" return "get_news"
# Registering functions to the router # Registering functions to the router
route_get_time = generate_route(get_time) route_get_time = generate_route(get_time)
route_get_news = generate_route(get_news) route_get_news = generate_route(get_news)
routes = [route_get_time, route_get_news] routes = [route_get_time, route_get_news]
router = create_router(routes) router = create_router(routes)
# Tools # Tools
tools = [get_time, get_news] tools = [get_time, get_news]
``` ```
%% Cell type:code id: tags: %% Cell type:code id: tags:
``` python ``` python
call(query="What is the time in Stockholm?", functions=tools, router=router) call(query="What is the time in Stockholm?", functions=tools, router=router)
call(query="What is the tech news in the Lithuania?", functions=tools, router=router) call(query="What is the tech news in the Lithuania?", functions=tools, router=router)
call(query="Hi!", functions=tools, router=router) call(query="Hi!", functions=tools, router=router)
``` ```
......
File added
This diff is collapsed.
[tool.poetry] [tool.poetry]
name = "semantic-router" name = "semantic-router"
version = "0.0.10" version = "0.0.11"
description = "Super fast semantic router for AI decision making" description = "Super fast semantic router for AI decision making"
authors = [ authors = [
"James Briggs <james@aurelio.ai>", "James Briggs <james@aurelio.ai>",
...@@ -12,7 +12,7 @@ authors = [ ...@@ -12,7 +12,7 @@ authors = [
readme = "README.md" readme = "README.md"
[tool.poetry.dependencies] [tool.poetry.dependencies]
python = "^3.10" python = "^3.9"
pydantic = "^1.8.2" pydantic = "^1.8.2"
openai = "^1.3.9" openai = "^1.3.9"
cohere = "^4.32" cohere = "^4.32"
......
...@@ -54,5 +54,5 @@ class OpenAIEncoder(BaseEncoder): ...@@ -54,5 +54,5 @@ class OpenAIEncoder(BaseEncoder):
): ):
raise ValueError(f"No embeddings returned. Error: {error_message}") raise ValueError(f"No embeddings returned. Error: {error_message}")
embeddings = [r.embedding for r in embeds.data] embeddings = [embeds_obj.embedding for embeds_obj in embeds.data]
return embeddings return embeddings
import pytest import pytest
from openai import OpenAIError from openai import OpenAIError
from openai.types import CreateEmbeddingResponse, Embedding from openai.types import CreateEmbeddingResponse, Embedding
from openai.types.create_embedding_response import Usage
from semantic_router.encoders import OpenAIEncoder from semantic_router.encoders import OpenAIEncoder
...@@ -41,6 +42,11 @@ class TestOpenAIEncoder: ...@@ -41,6 +42,11 @@ class TestOpenAIEncoder:
) )
def test_openai_encoder_call_success(self, openai_encoder, mocker): def test_openai_encoder_call_success(self, openai_encoder, mocker):
mock_embeddings = mocker.Mock()
mock_embeddings.data = [
Embedding(embedding=[0.1, 0.2], index=0, object="embedding")
]
mocker.patch("os.getenv", return_value="fake-api-key") mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch("time.sleep", return_value=None) # To speed up the test mocker.patch("time.sleep", return_value=None) # To speed up the test
...@@ -49,7 +55,7 @@ class TestOpenAIEncoder: ...@@ -49,7 +55,7 @@ class TestOpenAIEncoder:
mock_response = CreateEmbeddingResponse( mock_response = CreateEmbeddingResponse(
model="text-embedding-ada-002", model="text-embedding-ada-002",
object="list", object="list",
usage={"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 20}, usage=Usage(prompt_tokens=0, total_tokens=20),
data=[mock_embedding], data=[mock_embedding],
) )
...@@ -70,7 +76,7 @@ class TestOpenAIEncoder: ...@@ -70,7 +76,7 @@ class TestOpenAIEncoder:
) )
with pytest.raises(ValueError) as e: with pytest.raises(ValueError) as e:
openai_encoder(["test document"]) openai_encoder(["test document"])
assert "No embeddings returned. Error: Test error" in str(e.value) assert "No embeddings returned. Error" in str(e.value)
def test_openai_encoder_call_failure_non_openai_error(self, openai_encoder, mocker): def test_openai_encoder_call_failure_non_openai_error(self, openai_encoder, mocker):
mocker.patch("os.getenv", return_value="fake-api-key") mocker.patch("os.getenv", return_value="fake-api-key")
...@@ -86,6 +92,11 @@ class TestOpenAIEncoder: ...@@ -86,6 +92,11 @@ class TestOpenAIEncoder:
assert "OpenAI API call failed. Error: Non-OpenAIError" in str(e.value) assert "OpenAI API call failed. Error: Non-OpenAIError" in str(e.value)
def test_openai_encoder_call_successful_retry(self, openai_encoder, mocker): def test_openai_encoder_call_successful_retry(self, openai_encoder, mocker):
mock_embeddings = mocker.Mock()
mock_embeddings.data = [
Embedding(embedding=[0.1, 0.2], index=0, object="embedding")
]
mocker.patch("os.getenv", return_value="fake-api-key") mocker.patch("os.getenv", return_value="fake-api-key")
mocker.patch("time.sleep", return_value=None) # To speed up the test mocker.patch("time.sleep", return_value=None) # To speed up the test
...@@ -94,7 +105,7 @@ class TestOpenAIEncoder: ...@@ -94,7 +105,7 @@ class TestOpenAIEncoder:
mock_response = CreateEmbeddingResponse( mock_response = CreateEmbeddingResponse(
model="text-embedding-ada-002", model="text-embedding-ada-002",
object="list", object="list",
usage={"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 20}, usage=Usage(prompt_tokens=0, total_tokens=20),
data=[mock_embedding], data=[mock_embedding],
) )
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment