diff --git a/libs/.docker/.env.example b/libs/.docker/.env.example index ef80dc16c..0facc4684 100644 --- a/libs/.docker/.env.example +++ b/libs/.docker/.env.example @@ -75,6 +75,8 @@ SEGMENT_WRITE_KEY= STRIPE_SECRET_KEY= # Research tool TAVILY_API_KEY= +# Azure OpenAI embeddings +AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT= ############################### diff --git a/libs/.docker/docker-compose.yml b/libs/.docker/docker-compose.yml index 926336265..3a3f21a83 100644 --- a/libs/.docker/docker-compose.yml +++ b/libs/.docker/docker-compose.yml @@ -54,6 +54,7 @@ services: - AGENTOPS_API_KEY=${AGENTOPS_API_KEY} - AGENTOPS_ORG_KEY=${AGENTOPS_ORG_KEY} - TAVILY_API_KEY=${TAVILY_API_KEY} + - AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT=${AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT} ports: - 8080:8080 restart: unless-stopped diff --git a/libs/superagent/.env.example b/libs/superagent/.env.example index f8c5e7622..66d60fc15 100644 --- a/libs/superagent/.env.example +++ b/libs/superagent/.env.example @@ -67,3 +67,5 @@ STRIPE_SECRET_KEY= # Tools TAVILY_API_KEY= +# Optional for Azure Embeddings +AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT= \ No newline at end of file diff --git a/libs/superagent/app/agents/langchain.py b/libs/superagent/app/agents/langchain.py index 73db7a71c..b6f269794 100644 --- a/libs/superagent/app/agents/langchain.py +++ b/libs/superagent/app/agents/langchain.py @@ -86,6 +86,10 @@ async def _get_tools( if agent_datasource.datasource.vectorDb else None ), + # TODO: This will be removed in v0.3 + # This is for the users who wants to + # use Azure both for LLM and embeddings + "embeddings_model_provider": self.agent_config.llms[0].llm.provider, "query_type": "document", } if tool_type == DatasourceTool diff --git a/libs/superagent/app/agents/llm.py b/libs/superagent/app/agents/llm.py index de436185d..b87332d93 100644 --- a/libs/superagent/app/agents/llm.py +++ b/libs/superagent/app/agents/llm.py @@ -69,7 +69,6 @@ class CustomAgentExecutor: async def ainvoke(self, input, *_, **kwargs): function_calling_res = {} - print("agent_config.tools", agent_config, input) if len(agent_config.tools) > 0: function_calling = await FunctionCalling( enable_streaming=False, diff --git a/libs/superagent/app/api/agents.py b/libs/superagent/app/api/agents.py index 36ead8837..03bb55ae6 100644 --- a/libs/superagent/app/api/agents.py +++ b/libs/superagent/app/api/agents.py @@ -91,11 +91,11 @@ async def get_llm_or_raise(data: LLMPayload) -> LLM: where={"provider": provider, "apiUserId": data.user_id} ) - # if not llm: - # raise HTTPException( - # status_code=status.HTTP_400_BAD_REQUEST, - # detail="Please set an LLM first", - # ) + if not llm: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Please set an LLM first", + ) return llm diff --git a/libs/superagent/app/api/datasources.py b/libs/superagent/app/api/datasources.py index dd6d78812..8b63b0b39 100644 --- a/libs/superagent/app/api/datasources.py +++ b/libs/superagent/app/api/datasources.py @@ -8,6 +8,7 @@ from app.datasource.flow import delete_datasource, vectorize_datasource from app.models.request import Datasource as DatasourceRequest +from app.models.request import EmbeddingsModelProvider from app.models.response import ( Datasource as DatasourceResponse, ) @@ -58,7 +59,7 @@ async def create( data = await prisma.datasource.create( { "apiUserId": api_user.id, - **body.dict(), + **body.dict(exclude={"embeddingsModelProvider"}), } ) @@ -66,6 +67,7 @@ async def run_vectorize_flow( datasource: Datasource, options: Optional[dict], vector_db_provider: Optional[str], + embeddings_model_provider: EmbeddingsModelProvider, ): try: await vectorize_datasource( @@ -73,6 +75,7 @@ async def run_vectorize_flow( # vector db configurations (api key, index name etc.) options=options, vector_db_provider=vector_db_provider, + embeddings_model_provider=embeddings_model_provider, ) except Exception as flow_exception: await prisma.datasource.update( @@ -88,6 +91,7 @@ async def run_vectorize_flow( vector_db_provider=( vector_db.provider if vector_db is not None else None ), + embeddings_model_provider=body.embeddingsModelProvider, ) ) return {"success": True, "data": data} diff --git a/libs/superagent/app/datasource/flow.py b/libs/superagent/app/datasource/flow.py index 302ddd331..88626f6f7 100644 --- a/libs/superagent/app/datasource/flow.py +++ b/libs/superagent/app/datasource/flow.py @@ -6,6 +6,7 @@ from app.datasource.loader import DataLoader from app.datasource.types import VALID_UNSTRUCTURED_DATA_TYPES +from app.models.request import EmbeddingsModelProvider from app.utils.prisma import prisma from app.vectorstores.base import VectorStoreMain from prisma.enums import DatasourceStatus @@ -39,12 +40,17 @@ async def handle_datasources( @task async def vectorize( - datasource: Datasource, options: Optional[dict], vector_db_provider: Optional[str] + datasource: Datasource, + options: Optional[dict], + vector_db_provider: Optional[str], + embeddings_model_provider: EmbeddingsModelProvider, ) -> None: data = DataLoader(datasource=datasource).load() vector_store = VectorStoreMain( - options=options, vector_db_provider=vector_db_provider + options=options, + vector_db_provider=vector_db_provider, + embeddings_model_provider=embeddings_model_provider, ) vector_store.embed_documents(documents=data, datasource_id=datasource.id) @@ -76,13 +82,17 @@ async def process_datasource(datasource_id: str, agent_id: str): retries=0, ) async def vectorize_datasource( - datasource: Datasource, options: Optional[dict], vector_db_provider: Optional[str] + datasource: Datasource, + options: Optional[dict], + vector_db_provider: Optional[str], + embeddings_model_provider: EmbeddingsModelProvider, ) -> None: if datasource.type in VALID_UNSTRUCTURED_DATA_TYPES: await vectorize( datasource=datasource, options=options, vector_db_provider=vector_db_provider, + embeddings_model_provider=embeddings_model_provider, ) await prisma.datasource.update( where={"id": datasource.id}, data={"status": DatasourceStatus.DONE} diff --git a/libs/superagent/app/models/request.py b/libs/superagent/app/models/request.py index eb0923593..b89921dfb 100644 --- a/libs/superagent/app/models/request.py +++ b/libs/superagent/app/models/request.py @@ -1,3 +1,4 @@ +from enum import Enum from typing import Any, Dict, List, Optional from openai.types.beta.assistant_create_params import Tool as OpenAiAssistantTool @@ -71,6 +72,11 @@ class AgentInvoke(BaseModel): llm_params: Optional[LLMParams] +class EmbeddingsModelProvider(str, Enum): + OPENAI = "OPENAI" + AZURE_OPENAI = "AZURE_OPENAI" + + class Datasource(BaseModel): name: str description: Optional[str] @@ -79,6 +85,9 @@ class Datasource(BaseModel): url: Optional[str] metadata: Optional[Dict[Any, Any]] vectorDbId: Optional[str] + embeddingsModelProvider: Optional[ + EmbeddingsModelProvider + ] = EmbeddingsModelProvider.OPENAI class DatasourceUpdate(BaseModel): diff --git a/libs/superagent/app/tools/datasource.py b/libs/superagent/app/tools/datasource.py index 0d591dabd..3dfaa4b3a 100644 --- a/libs/superagent/app/tools/datasource.py +++ b/libs/superagent/app/tools/datasource.py @@ -77,6 +77,7 @@ def _run( vector_store = VectorStoreMain( options=self.metadata["options"], vector_db_provider=self.metadata["provider"], + embeddings_model_provider=self.metadata["embeddings_model_provider"], ) result = vector_store.query_documents( prompt=question, @@ -94,6 +95,7 @@ async def _arun( vector_store = VectorStoreMain( options=self.metadata["options"], vector_db_provider=self.metadata["provider"], + embeddings_model_provider=self.metadata["embeddings_model_provider"], ) result = vector_store.query_documents( prompt=question, diff --git a/libs/superagent/app/vectorstores/astra.py b/libs/superagent/app/vectorstores/astra.py index b145b8437..a1488aade 100644 --- a/libs/superagent/app/vectorstores/astra.py +++ b/libs/superagent/app/vectorstores/astra.py @@ -1,17 +1,17 @@ import logging -import os import uuid from typing import List, Literal, Optional import backoff from decouple import config from langchain.docstore.document import Document -from langchain.embeddings.openai import OpenAIEmbeddings # type: ignore from pydantic.dataclasses import dataclass +from app.models.request import EmbeddingsModelProvider from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase from app.vectorstores.astra_client import AstraClient, QueryResponse +from app.vectorstores.embeddings import get_embeddings_model_provider logger = logging.getLogger(__name__) @@ -40,6 +40,7 @@ class AstraVectorStore(VectorStoreBase): def __init__( self, options: dict, + embeddings_model_provider: EmbeddingsModelProvider, astra_id: str = None, astra_region: str = None, astra_application_token: str = None, @@ -91,9 +92,8 @@ def __init__( variables["ASTRA_DB_COLLECTION_NAME"], ) - self.embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - openai_api_key=os.getenv("OPENAI_API_KEY", ""), + self.embeddings = get_embeddings_model_provider( + embeddings_model_provider=embeddings_model_provider ) @backoff.on_exception(backoff.expo, Exception, max_tries=3) diff --git a/libs/superagent/app/vectorstores/base.py b/libs/superagent/app/vectorstores/base.py index fdf0e6926..ec77c27d5 100644 --- a/libs/superagent/app/vectorstores/base.py +++ b/libs/superagent/app/vectorstores/base.py @@ -4,6 +4,7 @@ from decouple import config from langchain.docstore.document import Document +from app.models.request import EmbeddingsModelProvider from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase from app.vectorstores.astra import AstraVectorStore @@ -28,7 +29,12 @@ class VectorStoreMain(VectorStoreBase): - def __init__(self, options: Optional[dict], vector_db_provider: Optional[str]): + def __init__( + self, + options: Optional[dict], + vector_db_provider: Optional[str], + embeddings_model_provider: EmbeddingsModelProvider, + ): """ Determine the vectorstore """ @@ -41,6 +47,7 @@ def __init__(self, options: Optional[dict], vector_db_provider: Optional[str]): VECTOR_DB_MAPPING.get(config("VECTORSTORE", None)), VectorDbProvider.PINECONE.value, ) + self.embeddings_model_provider = embeddings_model_provider self.instance = self.get_database() def get_database(self, index_name: Optional[str] = None) -> Any: @@ -84,7 +91,9 @@ def get_database(self, index_name: Optional[str] = None) -> Any: if index_name is None: index_name = index_names.get(self.vectorstore) return vectorstore_classes.get(self.vectorstore)( - index_name=index_name, options=self.options + index_name=index_name, + options=self.options, + embeddings_model_provider=self.embeddings_model_provider, ) def query( diff --git a/libs/superagent/app/vectorstores/embeddings.py b/libs/superagent/app/vectorstores/embeddings.py new file mode 100644 index 000000000..74050a46c --- /dev/null +++ b/libs/superagent/app/vectorstores/embeddings.py @@ -0,0 +1,28 @@ +from decouple import config +from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings + +from app.models.request import EmbeddingsModelProvider +from app.utils.helpers import get_first_non_null + + +def get_embeddings_model_provider(embeddings_model_provider: EmbeddingsModelProvider): + if embeddings_model_provider == EmbeddingsModelProvider.AZURE_OPENAI: + return AzureOpenAIEmbeddings( + azure_deployment=config("AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT"), + api_version=get_first_non_null( + config("AZURE_OPENAI_EMBEDDINGS_API_VERSION"), + config("AZURE_OPENAI_API_VERSION"), + ), + api_key=get_first_non_null( + config("AZURE_OPENAI_EMBEDDINGS_API_KEY"), + config("AZURE_OPENAI_API_KEY"), + ), + azure_endpoint=get_first_non_null( + config("AZURE_OPENAI_EMBEDDINGS_ENDPOINT"), + config("AZURE_OPENAI_ENDPOINT"), + ), + ) + else: + return OpenAIEmbeddings( + model="text-embedding-3-small", openai_api_key=config("OPENAI_API_KEY") + ) diff --git a/libs/superagent/app/vectorstores/pinecone.py b/libs/superagent/app/vectorstores/pinecone.py index f3cd045cb..5d4c39f95 100644 --- a/libs/superagent/app/vectorstores/pinecone.py +++ b/libs/superagent/app/vectorstores/pinecone.py @@ -6,12 +6,13 @@ import pinecone from decouple import config from langchain.docstore.document import Document -from langchain.embeddings.openai import OpenAIEmbeddings # type: ignore from pinecone.core.client.models import QueryResponse from pydantic.dataclasses import dataclass +from app.models.request import EmbeddingsModelProvider from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase +from app.vectorstores.embeddings import get_embeddings_model_provider logger = logging.getLogger(__name__) @@ -40,6 +41,7 @@ class PineconeVectorStore(VectorStoreBase): def __init__( self, options: dict, + embeddings_model_provider: EmbeddingsModelProvider, index_name: str = None, environment: str = None, pinecone_api_key: str = None, @@ -82,9 +84,7 @@ def __init__( self.index_name = variables["PINECONE_INDEX"] logger.info(f"Index name: {self.index_name}") self.index = pinecone.Index(self.index_name) - self.embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", openai_api_key=config("OPENAI_API_KEY") - ) # type: ignore + self.embeddings = get_embeddings_model_provider(embeddings_model_provider) @backoff.on_exception(backoff.expo, Exception, max_tries=3) def _embed_with_retry(self, texts): diff --git a/libs/superagent/app/vectorstores/qdrant.py b/libs/superagent/app/vectorstores/qdrant.py index e926c0ff7..c7623a361 100644 --- a/libs/superagent/app/vectorstores/qdrant.py +++ b/libs/superagent/app/vectorstores/qdrant.py @@ -1,16 +1,16 @@ import logging from typing import Literal -import openai from decouple import config from langchain.docstore.document import Document -from langchain.embeddings.openai import OpenAIEmbeddings # type: ignore from qdrant_client import QdrantClient, models from qdrant_client.http import models as rest from qdrant_client.http.models import PointStruct +from app.models.request import EmbeddingsModelProvider from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase +from app.vectorstores.embeddings import get_embeddings_model_provider logger = logging.getLogger(__name__) @@ -19,6 +19,7 @@ class QdrantVectorStore(VectorStoreBase): def __init__( self, options: dict, + embeddings_model_provider: EmbeddingsModelProvider, index_name: str = None, host: str = None, api_key: str = None, @@ -55,8 +56,8 @@ def __init__( url=variables["QDRANT_HOST"], api_key=variables["QDRANT_API_KEY"], ) - self.embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", openai_api_key=config("OPENAI_API_KEY") + self.embeddings = get_embeddings_model_provider( + embeddings_model_provider=embeddings_model_provider ) self.index_name = variables["QDRANT_INDEX"] @@ -78,13 +79,11 @@ def embed_documents(self, documents: list[Document], batch_size: int = 100) -> N i = 0 for document in documents: i += 1 - response = openai.embeddings.create( - input=document.page_content, model="text-embedding-3-small" - ) + response = self.embeddings.embed_documents([document.page_content]) points.append( PointStruct( id=i, - vector={"content": response.data[0].embedding}, + vector={"content": response[0]}, payload={"text": document.page_content, **document.metadata}, ) ) @@ -97,10 +96,8 @@ def query_documents( top_k: int | None, _query_type: Literal["document", "all"] = "document", ) -> list[str]: - response = openai.embeddings.create( - input=prompt, model="text-embedding-3-small" - ) - embeddings = response.data[0].embedding + response = self.embeddings.embed_documents([prompt]) + embeddings = response[0] search_result = self.client.search( collection_name=self.index_name, query_vector=("content", embeddings), diff --git a/libs/superagent/app/vectorstores/supabase.py b/libs/superagent/app/vectorstores/supabase.py index 584cc9b4c..2920fd025 100644 --- a/libs/superagent/app/vectorstores/supabase.py +++ b/libs/superagent/app/vectorstores/supabase.py @@ -6,9 +6,10 @@ import backoff from decouple import config from langchain.docstore.document import Document -from langchain.embeddings.openai import OpenAIEmbeddings # type: ignore from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase +from app.vectorstores.embeddings import get_embeddings_model_provider +from app.models.request import EmbeddingsModelProvider logger = logging.getLogger(__name__) @@ -17,6 +18,7 @@ class SupabaseVectorStore(VectorStoreBase): def __init__( self, options: dict, + embeddings_model_provider: EmbeddingsModelProvider, index_name: str = None, db_conn_url: str = None, url: str = None, @@ -45,9 +47,7 @@ def __init__( # create vector store client self.client = vecs.create_client(variables["SUPABASE_DB_URL"]) - self.embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", openai_api_key=config("OPENAI_API_KEY") - ) + self.embeddings = get_embeddings_model_provider(embeddings_model_provider) # create a collection named 'sentences' with 1536 dimensional vectors (default dimension for text-embedding-3-small) self.collection = self.client.get_or_create_collection( diff --git a/libs/superagent/app/vectorstores/weaviate.py b/libs/superagent/app/vectorstores/weaviate.py index 392d34873..e01cddc95 100644 --- a/libs/superagent/app/vectorstores/weaviate.py +++ b/libs/superagent/app/vectorstores/weaviate.py @@ -7,10 +7,11 @@ import weaviate from decouple import config from langchain.docstore.document import Document -from langchain.embeddings.openai import OpenAIEmbeddings # type: ignore from pydantic.dataclasses import dataclass from app.utils.helpers import get_first_non_null from app.vectorstores.abstract import VectorStoreBase +from app.vectorstores.embeddings import get_embeddings_model_provider +from app.models.request import EmbeddingsModelProvider logger = logging.getLogger(__name__) @@ -51,6 +52,7 @@ class WeaviateVectorStore(VectorStoreBase): def __init__( self, options: dict, + embeddings_model_provider: EmbeddingsModelProvider, index_name: str = None, api_key: str = None, url: str = None, @@ -85,8 +87,8 @@ def __init__( url=variables["WEAVIATE_URL"], auth_client_secret=auth, ) - self.embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", openai_api_key=config("OPENAI_API_KEY") + self.embeddings = get_embeddings_model_provider( + embeddings_model_provider=embeddings_model_provider ) self.index_name = variables["WEAVIATE_INDEX"] diff --git a/libs/ui/app/agents/[agentId]/add-datasource.tsx b/libs/ui/app/agents/[agentId]/add-datasource.tsx index db3a63d5a..9d000a8a9 100644 --- a/libs/ui/app/agents/[agentId]/add-datasource.tsx +++ b/libs/ui/app/agents/[agentId]/add-datasource.tsx @@ -2,6 +2,7 @@ import * as React from "react" import { useRouter } from "next/navigation" +import { LLMProvider } from "@/models/models" import { zodResolver } from "@hookform/resolvers/zod" import { createClientComponentClient } from "@supabase/auth-helpers-nextjs" import { useForm } from "react-hook-form" @@ -47,15 +48,19 @@ const formSchema = z.object({ metadata: z.any(), }) +interface AddDatasourceProps { + profile: any + agent: any + onSuccess: () => void + llmProvider: LLMProvider +} + function AddDatasource({ profile, agent, onSuccess, -}: { - profile: any - agent: any - onSuccess: () => void -}) { + llmProvider, +}: AddDatasourceProps) { const supabase = createClientComponentClient() const router = useRouter() const { toast } = useToast() @@ -90,6 +95,7 @@ function AddDatasource({ const { data: datasource } = await api.createDatasource({ ...values, vectorDbId: vectorDbs[0]?.id, + embeddingsModelProvider: getEmbeddingsModelProvider(llmProvider), }) await api.createAgentDatasource(agent.id, datasource.id) form.reset() @@ -106,6 +112,13 @@ function AddDatasource({ } } + function getEmbeddingsModelProvider(llmProvider: LLMProvider): LLMProvider { + if (llmProvider === LLMProvider.AZURE_OPENAI) + return LLMProvider.AZURE_OPENAI + + return LLMProvider.OPENAI + } + function mapMimeTypeToFileType(mimeType: string): string { const typeMapping: { [key: string]: string } = { "text/plain": "TXT", diff --git a/libs/ui/app/agents/[agentId]/settings.tsx b/libs/ui/app/agents/[agentId]/settings.tsx index 2fdfb3ac3..e30b5ca8c 100644 --- a/libs/ui/app/agents/[agentId]/settings.tsx +++ b/libs/ui/app/agents/[agentId]/settings.tsx @@ -2,6 +2,7 @@ import * as React from "react" import { useRouter } from "next/navigation" +import { LLMProvider } from "@/models/models" import { zodResolver } from "@hookform/resolvers/zod" import { useForm } from "react-hook-form" import * as z from "zod" @@ -77,6 +78,7 @@ export default function Settings({ tools, profile, }: SettingsProps) { + console.log("agent", agent) const api = new Api(profile.api_key) const router = useRouter() const { toast } = useToast() @@ -95,7 +97,7 @@ export default function Settings({ }, }) const avatar = form.watch("avatar") - const llms = form.watch("llms") + const currLlmProvider = form.watch("llms") as LLMProvider async function onSubmit(values: z.infer) { const { tools, datasources } = values @@ -140,9 +142,9 @@ export default function Settings({ (datasourceId) => api.deleteAgentDatasource(agent.id, datasourceId) ) - if (llms !== agent.llms?.[0]?.llm.provider) { + if (currLlmProvider !== agent.llms?.[0]?.llm.provider) { const configuredLLM = configuredLLMs.find( - (llm) => llm.provider === llms + (llm) => llm.provider === currLlmProvider ) if (configuredLLM) { @@ -278,7 +280,7 @@ export default function Settings({ {siteConfig.llms - .find((llm) => llm.id === llms) + .find((llm) => llm.id === currLlmProvider) ?.options.map((option) => ( { window.location.reload() }} diff --git a/libs/ui/models/models.ts b/libs/ui/models/models.ts index 54702a511..188166a3e 100644 --- a/libs/ui/models/models.ts +++ b/libs/ui/models/models.ts @@ -1,7 +1,7 @@ export enum LLMProvider { - OPENAI, - AZURE_OPENAI, - HUGGINGFACE, + OPENAI = "OPENAI", + AZURE_OPENAI = "AZURE_OPENAI", + HUGGINGFACE = "HUGGINGFACE", } export enum LLMModel {