Rename embedchain to mem0 and open sourcing code for long term memory (#1474)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Taranjeet Singh
co-authored by
Deshraj Yadav
Deshraj Yadav
parent
83e8c97295
commit
f842a92e25
@@ -0,0 +1,59 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
try:
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
except ImportError:
|
||||
raise ImportError("Please install the langchain-anthropic package by running `pip install langchain-anthropic`.")
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class AnthropicLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "ANTHROPIC_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the ANTHROPIC_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "anthropic/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["input_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["output_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["input_tokens"],
|
||||
"completion_tokens": token_info["output_tokens"],
|
||||
"total_tokens": token_info["input_tokens"] + token_info["output_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
api_key = config.api_key or os.getenv("ANTHROPIC_API_KEY")
|
||||
chat = ChatAnthropic(anthropic_api_key=api_key, temperature=config.temperature, model_name=config.model)
|
||||
|
||||
if config.max_tokens and config.max_tokens != 1000:
|
||||
logger.warning("Config option `max_tokens` is not supported by this model.")
|
||||
|
||||
messages = BaseLlm._get_messages(prompt, system_prompt=config.system_prompt)
|
||||
|
||||
chat_response = chat.invoke(messages)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.llms import Bedrock
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class AWSBedrockLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config)
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> str:
|
||||
response = self._get_answer(prompt, self.config)
|
||||
return response
|
||||
|
||||
def _get_answer(self, prompt: str, config: BaseLlmConfig) -> str:
|
||||
try:
|
||||
import boto3
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for AWSBedrock are not installed."
|
||||
"Please install with `pip install boto3==1.34.20`"
|
||||
) from None
|
||||
|
||||
self.boto_client = boto3.client("bedrock-runtime", "us-west-2" or os.environ.get("AWS_REGION"))
|
||||
|
||||
kwargs = {
|
||||
"model_id": config.model or "amazon.titan-text-express-v1",
|
||||
"client": self.boto_client,
|
||||
"model_kwargs": config.model_kwargs
|
||||
or {
|
||||
"temperature": config.temperature,
|
||||
},
|
||||
}
|
||||
|
||||
if config.stream:
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
|
||||
callbacks = [StreamingStdOutCallbackHandler()]
|
||||
llm = Bedrock(**kwargs, streaming=config.stream, callbacks=callbacks)
|
||||
else:
|
||||
llm = Bedrock(**kwargs)
|
||||
|
||||
return llm.invoke(prompt)
|
||||
@@ -0,0 +1,40 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class AzureOpenAILlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
return self._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
from langchain_openai import AzureChatOpenAI
|
||||
|
||||
if not config.deployment_name:
|
||||
raise ValueError("Deployment name must be provided for Azure OpenAI")
|
||||
|
||||
chat = AzureChatOpenAI(
|
||||
deployment_name=config.deployment_name,
|
||||
openai_api_version=str(config.api_version) if config.api_version else "2024-02-01",
|
||||
model_name=config.model or "gpt-3.5-turbo",
|
||||
temperature=config.temperature,
|
||||
max_tokens=config.max_tokens,
|
||||
streaming=config.stream,
|
||||
)
|
||||
|
||||
if config.top_p and config.top_p != 1:
|
||||
logger.warning("Config option `top_p` is not supported by this model.")
|
||||
|
||||
messages = BaseLlm._get_messages(prompt, system_prompt=config.system_prompt)
|
||||
|
||||
return chat.invoke(messages).content
|
||||
@@ -0,0 +1,342 @@
|
||||
import logging
|
||||
from collections.abc import Generator
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain.schema import BaseMessage as LCBaseMessage
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.config.llm.base import (
|
||||
DEFAULT_PROMPT,
|
||||
DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE,
|
||||
DEFAULT_PROMPT_WITH_MEM0_MEMORY_TEMPLATE,
|
||||
DOCS_SITE_PROMPT_TEMPLATE,
|
||||
)
|
||||
from embedchain.helpers.json_serializable import JSONSerializable
|
||||
from embedchain.memory.base import ChatHistory
|
||||
from embedchain.memory.message import ChatMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseLlm(JSONSerializable):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
"""Initialize a base LLM class
|
||||
|
||||
:param config: LLM configuration option class, defaults to None
|
||||
:type config: Optional[BaseLlmConfig], optional
|
||||
"""
|
||||
if config is None:
|
||||
self.config = BaseLlmConfig()
|
||||
else:
|
||||
self.config = config
|
||||
|
||||
self.memory = ChatHistory()
|
||||
self.is_docs_site_instance = False
|
||||
self.history: Any = None
|
||||
|
||||
def get_llm_model_answer(self):
|
||||
"""
|
||||
Usually implemented by child class
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def set_history(self, history: Any):
|
||||
"""
|
||||
Provide your own history.
|
||||
Especially interesting for the query method, which does not internally manage conversation history.
|
||||
|
||||
:param history: History to set
|
||||
:type history: Any
|
||||
"""
|
||||
self.history = history
|
||||
|
||||
def update_history(self, app_id: str, session_id: str = "default"):
|
||||
"""Update class history attribute with history in memory (for chat method)"""
|
||||
chat_history = self.memory.get(app_id=app_id, session_id=session_id, num_rounds=10)
|
||||
self.set_history([str(history) for history in chat_history])
|
||||
|
||||
def add_history(
|
||||
self,
|
||||
app_id: str,
|
||||
question: str,
|
||||
answer: str,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
session_id: str = "default",
|
||||
):
|
||||
chat_message = ChatMessage()
|
||||
chat_message.add_user_message(question, metadata=metadata)
|
||||
chat_message.add_ai_message(answer, metadata=metadata)
|
||||
self.memory.add(app_id=app_id, chat_message=chat_message, session_id=session_id)
|
||||
self.update_history(app_id=app_id, session_id=session_id)
|
||||
|
||||
def _format_history(self) -> str:
|
||||
"""Format history to be used in prompt
|
||||
|
||||
:return: Formatted history
|
||||
:rtype: str
|
||||
"""
|
||||
return "\n".join(self.history)
|
||||
|
||||
def _format_memories(self, memories: list[dict]) -> str:
|
||||
"""Format memories to be used in prompt
|
||||
|
||||
:param memories: Memories to format
|
||||
:type memories: list[dict]
|
||||
:return: Formatted memories
|
||||
:rtype: str
|
||||
"""
|
||||
return "\n".join([memory["text"] for memory in memories])
|
||||
|
||||
def generate_prompt(self, input_query: str, contexts: list[str], **kwargs: dict[str, Any]) -> str:
|
||||
"""
|
||||
Generates a prompt based on the given query and context, ready to be
|
||||
passed to an LLM
|
||||
|
||||
:param input_query: The query to use.
|
||||
:type input_query: str
|
||||
:param contexts: List of similar documents to the query used as context.
|
||||
:type contexts: list[str]
|
||||
:return: The prompt
|
||||
:rtype: str
|
||||
"""
|
||||
context_string = " | ".join(contexts)
|
||||
web_search_result = kwargs.get("web_search_result", "")
|
||||
memories = kwargs.get("memories", None)
|
||||
if web_search_result:
|
||||
context_string = self._append_search_and_context(context_string, web_search_result)
|
||||
|
||||
prompt_contains_history = self.config._validate_prompt_history(self.config.prompt)
|
||||
if prompt_contains_history:
|
||||
prompt = self.config.prompt.substitute(
|
||||
context=context_string, query=input_query, history=self._format_history() or "No history"
|
||||
)
|
||||
elif self.history and not prompt_contains_history:
|
||||
# History is present, but not included in the prompt.
|
||||
# check if it's the default prompt without history
|
||||
if (
|
||||
not self.config._validate_prompt_history(self.config.prompt)
|
||||
and self.config.prompt.template == DEFAULT_PROMPT
|
||||
):
|
||||
if memories:
|
||||
# swap in the template with Mem0 memory template
|
||||
prompt = DEFAULT_PROMPT_WITH_MEM0_MEMORY_TEMPLATE.substitute(
|
||||
context=context_string,
|
||||
query=input_query,
|
||||
history=self._format_history(),
|
||||
memories=self._format_memories(memories),
|
||||
)
|
||||
else:
|
||||
# swap in the template with history
|
||||
prompt = DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE.substitute(
|
||||
context=context_string, query=input_query, history=self._format_history()
|
||||
)
|
||||
else:
|
||||
# If we can't swap in the default, we still proceed but tell users that the history is ignored.
|
||||
logger.warning(
|
||||
"Your bot contains a history, but prompt does not include `$history` key. History is ignored."
|
||||
)
|
||||
prompt = self.config.prompt.substitute(context=context_string, query=input_query)
|
||||
else:
|
||||
# basic use case, no history.
|
||||
prompt = self.config.prompt.substitute(context=context_string, query=input_query)
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
def _append_search_and_context(context: str, web_search_result: str) -> str:
|
||||
"""Append web search context to existing context
|
||||
|
||||
:param context: Existing context
|
||||
:type context: str
|
||||
:param web_search_result: Web search result
|
||||
:type web_search_result: str
|
||||
:return: Concatenated web search result
|
||||
:rtype: str
|
||||
"""
|
||||
return f"{context}\nWeb Search Result: {web_search_result}"
|
||||
|
||||
def get_answer_from_llm(self, prompt: str):
|
||||
"""
|
||||
Gets an answer based on the given query and context by passing it
|
||||
to an LLM.
|
||||
|
||||
:param prompt: Gets an answer based on the given query and context by passing it to an LLM.
|
||||
:type prompt: str
|
||||
:return: The answer.
|
||||
:rtype: _type_
|
||||
"""
|
||||
return self.get_llm_model_answer(prompt)
|
||||
|
||||
@staticmethod
|
||||
def access_search_and_get_results(input_query: str):
|
||||
"""
|
||||
Search the internet for additional context
|
||||
|
||||
:param input_query: search query
|
||||
:type input_query: str
|
||||
:return: Search results
|
||||
:rtype: Unknown
|
||||
"""
|
||||
try:
|
||||
from langchain.tools import DuckDuckGoSearchRun
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Searching requires extra dependencies. Install with `pip install duckduckgo-search==6.1.5`"
|
||||
) from None
|
||||
search = DuckDuckGoSearchRun()
|
||||
logger.info(f"Access search to get answers for {input_query}")
|
||||
return search.run(input_query)
|
||||
|
||||
@staticmethod
|
||||
def _stream_response(answer: Any, token_info: Optional[dict[str, Any]] = None) -> Generator[Any, Any, None]:
|
||||
"""Generator to be used as streaming response
|
||||
|
||||
:param answer: Answer chunk from llm
|
||||
:type answer: Any
|
||||
:yield: Answer chunk from llm
|
||||
:rtype: Generator[Any, Any, None]
|
||||
"""
|
||||
streamed_answer = ""
|
||||
for chunk in answer:
|
||||
streamed_answer = streamed_answer + chunk
|
||||
yield chunk
|
||||
logger.info(f"Answer: {streamed_answer}")
|
||||
if token_info:
|
||||
logger.info(f"Token Info: {token_info}")
|
||||
|
||||
def query(self, input_query: str, contexts: list[str], config: BaseLlmConfig = None, dry_run=False, memories=None):
|
||||
"""
|
||||
Queries the vector database based on the given input query.
|
||||
Gets relevant doc based on the query and then passes it to an
|
||||
LLM as context to get the answer.
|
||||
|
||||
:param input_query: The query to use.
|
||||
:type input_query: str
|
||||
:param contexts: Embeddings retrieved from the database to be used as context.
|
||||
:type contexts: list[str]
|
||||
:param config: The `BaseLlmConfig` instance to use as configuration options. This is used for one method call.
|
||||
To persistently use a config, declare it during app init., defaults to None
|
||||
:type config: Optional[BaseLlmConfig], optional
|
||||
:param dry_run: A dry run does everything except send the resulting prompt to
|
||||
the LLM. The purpose is to test the prompt, not the response., defaults to False
|
||||
:type dry_run: bool, optional
|
||||
:return: The answer to the query or the dry run result
|
||||
:rtype: str
|
||||
"""
|
||||
try:
|
||||
if config:
|
||||
# A config instance passed to this method will only be applied temporarily, for one call.
|
||||
# So we will save the previous config and restore it at the end of the execution.
|
||||
# For this we use the serializer.
|
||||
prev_config = self.config.serialize()
|
||||
self.config = config
|
||||
|
||||
if config is not None and config.query_type == "Images":
|
||||
return contexts
|
||||
|
||||
if self.is_docs_site_instance:
|
||||
self.config.prompt = DOCS_SITE_PROMPT_TEMPLATE
|
||||
self.config.number_documents = 5
|
||||
k = {}
|
||||
if self.config.online:
|
||||
k["web_search_result"] = self.access_search_and_get_results(input_query)
|
||||
k["memories"] = memories
|
||||
prompt = self.generate_prompt(input_query, contexts, **k)
|
||||
logger.info(f"Prompt: {prompt}")
|
||||
if dry_run:
|
||||
return prompt
|
||||
|
||||
if self.config.token_usage:
|
||||
answer, token_info = self.get_answer_from_llm(prompt)
|
||||
else:
|
||||
answer = self.get_answer_from_llm(prompt)
|
||||
if isinstance(answer, str):
|
||||
logger.info(f"Answer: {answer}")
|
||||
if self.config.token_usage:
|
||||
return answer, token_info
|
||||
return answer
|
||||
else:
|
||||
if self.config.token_usage:
|
||||
return self._stream_response(answer, token_info)
|
||||
return self._stream_response(answer)
|
||||
finally:
|
||||
if config:
|
||||
# Restore previous config
|
||||
self.config: BaseLlmConfig = BaseLlmConfig.deserialize(prev_config)
|
||||
|
||||
def chat(
|
||||
self, input_query: str, contexts: list[str], config: BaseLlmConfig = None, dry_run=False, session_id: str = None
|
||||
):
|
||||
"""
|
||||
Queries the vector database on the given input query.
|
||||
Gets relevant doc based on the query and then passes it to an
|
||||
LLM as context to get the answer.
|
||||
|
||||
Maintains the whole conversation in memory.
|
||||
|
||||
:param input_query: The query to use.
|
||||
:type input_query: str
|
||||
:param contexts: Embeddings retrieved from the database to be used as context.
|
||||
:type contexts: list[str]
|
||||
:param config: The `BaseLlmConfig` instance to use as configuration options. This is used for one method call.
|
||||
To persistently use a config, declare it during app init., defaults to None
|
||||
:type config: Optional[BaseLlmConfig], optional
|
||||
:param dry_run: A dry run does everything except send the resulting prompt to
|
||||
the LLM. The purpose is to test the prompt, not the response., defaults to False
|
||||
:type dry_run: bool, optional
|
||||
:param session_id: Session ID to use for the conversation, defaults to None
|
||||
:type session_id: str, optional
|
||||
:return: The answer to the query or the dry run result
|
||||
:rtype: str
|
||||
"""
|
||||
try:
|
||||
if config:
|
||||
# A config instance passed to this method will only be applied temporarily, for one call.
|
||||
# So we will save the previous config and restore it at the end of the execution.
|
||||
# For this we use the serializer.
|
||||
prev_config = self.config.serialize()
|
||||
self.config = config
|
||||
|
||||
if self.is_docs_site_instance:
|
||||
self.config.prompt = DOCS_SITE_PROMPT_TEMPLATE
|
||||
self.config.number_documents = 5
|
||||
k = {}
|
||||
if self.config.online:
|
||||
k["web_search_result"] = self.access_search_and_get_results(input_query)
|
||||
|
||||
prompt = self.generate_prompt(input_query, contexts, **k)
|
||||
logger.info(f"Prompt: {prompt}")
|
||||
|
||||
if dry_run:
|
||||
return prompt
|
||||
|
||||
answer, token_info = self.get_answer_from_llm(prompt)
|
||||
if isinstance(answer, str):
|
||||
logger.info(f"Answer: {answer}")
|
||||
return answer, token_info
|
||||
else:
|
||||
# this is a streamed response and needs to be handled differently.
|
||||
return self._stream_response(answer, token_info)
|
||||
finally:
|
||||
if config:
|
||||
# Restore previous config
|
||||
self.config: BaseLlmConfig = BaseLlmConfig.deserialize(prev_config)
|
||||
|
||||
@staticmethod
|
||||
def _get_messages(prompt: str, system_prompt: Optional[str] = None) -> list[LCBaseMessage]:
|
||||
"""
|
||||
Construct a list of langchain messages
|
||||
|
||||
:param prompt: User prompt
|
||||
:type prompt: str
|
||||
:param system_prompt: System prompt, defaults to None
|
||||
:type system_prompt: Optional[str], optional
|
||||
:return: List of messages
|
||||
:rtype: list[BaseMessage]
|
||||
"""
|
||||
from langchain.schema import HumanMessage, SystemMessage
|
||||
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append(SystemMessage(content=system_prompt))
|
||||
messages.append(HumanMessage(content=prompt))
|
||||
return messages
|
||||
@@ -0,0 +1,47 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class ClarifaiLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "CLARIFAI_PAT" not in os.environ:
|
||||
raise ValueError("Please set the CLARIFAI_PAT environment variable.")
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
return self._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
try:
|
||||
from clarifai.client.model import Model
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for Clarifai are not installed."
|
||||
"Please install with `pip install clarifai==10.0.1`"
|
||||
) from None
|
||||
|
||||
model_name = config.model
|
||||
logging.info(f"Using clarifai LLM model: {model_name}")
|
||||
api_key = config.api_key or os.getenv("CLARIFAI_PAT")
|
||||
model = Model(url=model_name, pat=api_key)
|
||||
params = config.model_kwargs
|
||||
|
||||
try:
|
||||
(params := {}) if config.model_kwargs is None else config.model_kwargs
|
||||
predict_response = model.predict_by_bytes(
|
||||
bytes(prompt, "utf-8"),
|
||||
input_type="text",
|
||||
inference_params=params,
|
||||
)
|
||||
text = predict_response.outputs[0].data.text.raw
|
||||
return text
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Predict failed, exception: {e}")
|
||||
@@ -0,0 +1,66 @@
|
||||
import importlib
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_cohere import ChatCohere
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class CohereLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
try:
|
||||
importlib.import_module("cohere")
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for Cohere are not installed."
|
||||
"Please install with `pip install langchain_cohere==1.16.0`"
|
||||
) from None
|
||||
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "COHERE_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the COHERE_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.system_prompt:
|
||||
raise ValueError("CohereLlm does not support `system_prompt`")
|
||||
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "cohere/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["input_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["output_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["input_tokens"],
|
||||
"completion_tokens": token_info["output_tokens"],
|
||||
"total_tokens": token_info["input_tokens"] + token_info["output_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
api_key = config.api_key or os.environ["COHERE_API_KEY"]
|
||||
kwargs = {
|
||||
"model_name": config.model or "command-r",
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"together_api_key": api_key,
|
||||
}
|
||||
|
||||
chat = ChatCohere(**kwargs)
|
||||
chat_response = chat.invoke(prompt)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_count"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,62 @@
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
except ImportError:
|
||||
raise ImportError("GoogleLlm requires extra dependencies. Install with `pip install google-generativeai`") from None
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class GoogleLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config)
|
||||
if not self.config.api_key and "GOOGLE_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the GOOGLE_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
|
||||
genai.configure(api_key=api_key)
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
if self.config.system_prompt:
|
||||
raise ValueError("GoogleLlm does not support `system_prompt`")
|
||||
response = self._get_answer(prompt)
|
||||
return response
|
||||
|
||||
def _get_answer(self, prompt: str) -> Union[str, Generator[Any, Any, None]]:
|
||||
model_name = self.config.model or "gemini-pro"
|
||||
logger.info(f"Using Google LLM model: {model_name}")
|
||||
model = genai.GenerativeModel(model_name=model_name)
|
||||
|
||||
generation_config_params = {
|
||||
"candidate_count": 1,
|
||||
"max_output_tokens": self.config.max_tokens,
|
||||
"temperature": self.config.temperature or 0.5,
|
||||
}
|
||||
|
||||
if 0.0 <= self.config.top_p <= 1.0:
|
||||
generation_config_params["top_p"] = self.config.top_p
|
||||
else:
|
||||
raise ValueError("`top_p` must be > 0.0 and < 1.0")
|
||||
|
||||
generation_config = genai.types.GenerationConfig(**generation_config_params)
|
||||
|
||||
response = model.generate_content(
|
||||
prompt,
|
||||
generation_config=generation_config,
|
||||
stream=self.config.stream,
|
||||
)
|
||||
if self.config.stream:
|
||||
# TODO: Implement streaming
|
||||
response.resolve()
|
||||
return response.text
|
||||
else:
|
||||
return response.text
|
||||
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
|
||||
from langchain.callbacks.stdout import StdOutCallbackHandler
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class GPT4ALLLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if self.config.model is None:
|
||||
self.config.model = "orca-mini-3b-gguf2-q4_0.gguf"
|
||||
self.instance = GPT4ALLLlm._get_instance(self.config.model)
|
||||
self.instance.streaming = self.config.stream
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
return self._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_instance(model):
|
||||
try:
|
||||
from langchain_community.llms.gpt4all import \
|
||||
GPT4All as LangchainGPT4All
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The GPT4All python package is not installed. Please install it with `pip install --upgrade embedchain[opensource]`" # noqa E501
|
||||
) from None
|
||||
|
||||
model_path = Path(model).expanduser()
|
||||
if os.path.isabs(model_path):
|
||||
if os.path.exists(model_path):
|
||||
return LangchainGPT4All(model=str(model_path))
|
||||
else:
|
||||
raise ValueError(f"Model does not exist at {model_path=}")
|
||||
else:
|
||||
return LangchainGPT4All(model=model, allow_download=True)
|
||||
|
||||
def _get_answer(self, prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
|
||||
if config.model and config.model != self.config.model:
|
||||
raise RuntimeError(
|
||||
"GPT4ALLLlm does not support switching models at runtime. Please create a new app instance."
|
||||
)
|
||||
|
||||
messages = []
|
||||
if config.system_prompt:
|
||||
messages.append(config.system_prompt)
|
||||
messages.append(prompt)
|
||||
kwargs = {
|
||||
"temp": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
}
|
||||
if config.top_p:
|
||||
kwargs["top_p"] = config.top_p
|
||||
|
||||
callbacks = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
|
||||
|
||||
response = self.instance.generate(prompts=messages, callbacks=callbacks, **kwargs)
|
||||
answer = ""
|
||||
for generations in response.generations:
|
||||
answer += " ".join(map(lambda generation: generation.text, generations))
|
||||
return answer
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
from langchain.schema import HumanMessage, SystemMessage
|
||||
|
||||
try:
|
||||
from langchain_groq import ChatGroq
|
||||
except ImportError:
|
||||
raise ImportError("Groq requires extra dependencies. Install with `pip install langchain-groq`") from None
|
||||
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class GroqLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "GROQ_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the GROQ_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "groq/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["completion_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["prompt_tokens"],
|
||||
"completion_tokens": token_info["completion_tokens"],
|
||||
"total_tokens": token_info["prompt_tokens"] + token_info["completion_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
def _get_answer(self, prompt: str, config: BaseLlmConfig) -> str:
|
||||
messages = []
|
||||
if config.system_prompt:
|
||||
messages.append(SystemMessage(content=config.system_prompt))
|
||||
messages.append(HumanMessage(content=prompt))
|
||||
api_key = config.api_key or os.environ["GROQ_API_KEY"]
|
||||
kwargs = {
|
||||
"model_name": config.model or "mixtral-8x7b-32768",
|
||||
"temperature": config.temperature,
|
||||
"groq_api_key": api_key,
|
||||
}
|
||||
if config.stream:
|
||||
callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
|
||||
chat = ChatGroq(**kwargs, streaming=config.stream, callbacks=callbacks, api_key=api_key)
|
||||
else:
|
||||
chat = ChatGroq(**kwargs)
|
||||
|
||||
chat_response = chat.invoke(prompt)
|
||||
if self.config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,99 @@
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint
|
||||
from langchain_community.llms.huggingface_hub import HuggingFaceHub
|
||||
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class HuggingFaceLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
try:
|
||||
importlib.import_module("huggingface_hub")
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for HuggingFaceHub are not installed."
|
||||
"Please install with `pip install huggingface-hub==0.23.0`"
|
||||
) from None
|
||||
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "HUGGINGFACE_ACCESS_TOKEN" not in os.environ:
|
||||
raise ValueError("Please set the HUGGINGFACE_ACCESS_TOKEN environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
if self.config.system_prompt:
|
||||
raise ValueError("HuggingFaceLlm does not support `system_prompt`")
|
||||
return HuggingFaceLlm._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
# If the user wants to run the model locally, they can do so by setting the `local` flag to True
|
||||
if config.model and config.local:
|
||||
return HuggingFaceLlm._from_pipeline(prompt=prompt, config=config)
|
||||
elif config.model:
|
||||
return HuggingFaceLlm._from_model(prompt=prompt, config=config)
|
||||
elif config.endpoint:
|
||||
return HuggingFaceLlm._from_endpoint(prompt=prompt, config=config)
|
||||
else:
|
||||
raise ValueError("Either `model` or `endpoint` must be set in config")
|
||||
|
||||
@staticmethod
|
||||
def _from_model(prompt: str, config: BaseLlmConfig) -> str:
|
||||
model_kwargs = {
|
||||
"temperature": config.temperature or 0.1,
|
||||
"max_new_tokens": config.max_tokens,
|
||||
}
|
||||
|
||||
if 0.0 < config.top_p < 1.0:
|
||||
model_kwargs["top_p"] = config.top_p
|
||||
else:
|
||||
raise ValueError("`top_p` must be > 0.0 and < 1.0")
|
||||
|
||||
model = config.model
|
||||
api_key = config.api_key or os.getenv("HUGGINGFACE_ACCESS_TOKEN")
|
||||
logger.info(f"Using HuggingFaceHub with model {model}")
|
||||
llm = HuggingFaceHub(
|
||||
huggingfacehub_api_token=api_key,
|
||||
repo_id=model,
|
||||
model_kwargs=model_kwargs,
|
||||
)
|
||||
return llm.invoke(prompt)
|
||||
|
||||
@staticmethod
|
||||
def _from_endpoint(prompt: str, config: BaseLlmConfig) -> str:
|
||||
api_key = config.api_key or os.getenv("HUGGINGFACE_ACCESS_TOKEN")
|
||||
llm = HuggingFaceEndpoint(
|
||||
huggingfacehub_api_token=api_key,
|
||||
endpoint_url=config.endpoint,
|
||||
task="text-generation",
|
||||
model_kwargs=config.model_kwargs,
|
||||
)
|
||||
return llm.invoke(prompt)
|
||||
|
||||
@staticmethod
|
||||
def _from_pipeline(prompt: str, config: BaseLlmConfig) -> str:
|
||||
model_kwargs = {
|
||||
"temperature": config.temperature or 0.1,
|
||||
"max_new_tokens": config.max_tokens,
|
||||
}
|
||||
|
||||
if 0.0 < config.top_p < 1.0:
|
||||
model_kwargs["top_p"] = config.top_p
|
||||
else:
|
||||
raise ValueError("`top_p` must be > 0.0 and < 1.0")
|
||||
|
||||
llm = HuggingFacePipeline.from_model_id(
|
||||
model_id=config.model,
|
||||
task="text-generation",
|
||||
pipeline_kwargs=model_kwargs,
|
||||
)
|
||||
return llm.invoke(prompt)
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain.schema import HumanMessage, SystemMessage
|
||||
from langchain_community.chat_models import JinaChat
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class JinaLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "JINACHAT_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the JINACHAT_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
response = JinaLlm._get_answer(prompt, self.config)
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
messages = []
|
||||
if config.system_prompt:
|
||||
messages.append(SystemMessage(content=config.system_prompt))
|
||||
messages.append(HumanMessage(content=prompt))
|
||||
kwargs = {
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"jinachat_api_key": config.api_key or os.environ["JINACHAT_API_KEY"],
|
||||
"model_kwargs": {},
|
||||
}
|
||||
if config.top_p:
|
||||
kwargs["model_kwargs"]["top_p"] = config.top_p
|
||||
if config.stream:
|
||||
from langchain.callbacks.streaming_stdout import \
|
||||
StreamingStdOutCallbackHandler
|
||||
|
||||
chat = JinaChat(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()])
|
||||
else:
|
||||
chat = JinaChat(**kwargs)
|
||||
return chat(messages).content
|
||||
@@ -0,0 +1,53 @@
|
||||
import importlib
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from langchain_community.llms.replicate import Replicate
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class Llama2Llm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
try:
|
||||
importlib.import_module("replicate")
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for Llama2 are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[llama2]"`'
|
||||
) from None
|
||||
|
||||
# Set default config values specific to this llm
|
||||
if not config:
|
||||
config = BaseLlmConfig()
|
||||
# Add variables to this block that have a default value in the parent class
|
||||
config.max_tokens = 500
|
||||
config.temperature = 0.75
|
||||
# Add variables that are `none` by default to this block.
|
||||
if not config.model:
|
||||
config.model = (
|
||||
"a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5"
|
||||
)
|
||||
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "REPLICATE_API_TOKEN" not in os.environ:
|
||||
raise ValueError("Please set the REPLICATE_API_TOKEN environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
# TODO: Move the model and other inputs into config
|
||||
if self.config.system_prompt:
|
||||
raise ValueError("Llama2 does not support `system_prompt`")
|
||||
api_key = self.config.api_key or os.getenv("REPLICATE_API_TOKEN")
|
||||
llm = Replicate(
|
||||
model=self.config.model,
|
||||
replicate_api_token=api_key,
|
||||
input={
|
||||
"temperature": self.config.temperature,
|
||||
"max_length": self.config.max_tokens,
|
||||
"top_p": self.config.top_p,
|
||||
},
|
||||
)
|
||||
return llm.invoke(prompt)
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class MistralAILlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config)
|
||||
if not self.config.api_key and "MISTRAL_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the MISTRAL_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "mistralai/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["completion_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["prompt_tokens"],
|
||||
"completion_tokens": token_info["completion_tokens"],
|
||||
"total_tokens": token_info["prompt_tokens"] + token_info["completion_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig):
|
||||
try:
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_mistralai.chat_models import ChatMistralAI
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for MistralAI are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[mistralai]"`'
|
||||
) from None
|
||||
|
||||
api_key = config.api_key or os.getenv("MISTRAL_API_KEY")
|
||||
client = ChatMistralAI(mistral_api_key=api_key)
|
||||
messages = []
|
||||
if config.system_prompt:
|
||||
messages.append(SystemMessage(content=config.system_prompt))
|
||||
messages.append(HumanMessage(content=prompt))
|
||||
kwargs = {
|
||||
"model": config.model or "mistral-tiny",
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"top_p": config.top_p,
|
||||
}
|
||||
|
||||
# TODO: Add support for streaming
|
||||
if config.stream:
|
||||
answer = ""
|
||||
for chunk in client.stream(**kwargs, input=messages):
|
||||
answer += chunk.content
|
||||
return answer
|
||||
else:
|
||||
chat_response = client.invoke(**kwargs, input=messages)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from langchain.callbacks.manager import CallbackManager
|
||||
from langchain.callbacks.stdout import StdOutCallbackHandler
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
|
||||
try:
|
||||
from langchain_nvidia_ai_endpoints import ChatNVIDIA
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"NVIDIA AI endpoints requires extra dependencies. Install with `pip install langchain-nvidia-ai-endpoints`"
|
||||
) from None
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class NvidiaLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "NVIDIA_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the NVIDIA_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "nvidia/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["input_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["output_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["input_tokens"],
|
||||
"completion_tokens": token_info["output_tokens"],
|
||||
"total_tokens": token_info["input_tokens"] + token_info["output_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
|
||||
callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
|
||||
model_kwargs = config.model_kwargs or {}
|
||||
labels = model_kwargs.get("labels", None)
|
||||
params = {"model": config.model, "nvidia_api_key": config.api_key or os.getenv("NVIDIA_API_KEY")}
|
||||
if config.system_prompt:
|
||||
params["system_prompt"] = config.system_prompt
|
||||
if config.temperature:
|
||||
params["temperature"] = config.temperature
|
||||
if config.top_p:
|
||||
params["top_p"] = config.top_p
|
||||
if labels:
|
||||
params["labels"] = labels
|
||||
llm = ChatNVIDIA(**params, callback_manager=CallbackManager(callback_manager))
|
||||
chat_response = llm.invoke(prompt) if labels is None else llm.invoke(prompt, labels=labels)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,54 @@
|
||||
import logging
|
||||
from collections.abc import Iterable
|
||||
from typing import Optional, Union
|
||||
|
||||
from langchain.callbacks.manager import CallbackManager
|
||||
from langchain.callbacks.stdout import StdOutCallbackHandler
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
from langchain_community.llms.ollama import Ollama
|
||||
|
||||
try:
|
||||
from ollama import Client
|
||||
except ImportError:
|
||||
raise ImportError("Ollama requires extra dependencies. Install with `pip install ollama`") from None
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class OllamaLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if self.config.model is None:
|
||||
self.config.model = "llama2"
|
||||
|
||||
client = Client(host=config.base_url)
|
||||
local_models = client.list()["models"]
|
||||
if not any(model.get("name") == self.config.model for model in local_models):
|
||||
logger.info(f"Pulling {self.config.model} from Ollama!")
|
||||
client.pull(self.config.model)
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
return self._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
|
||||
if config.stream:
|
||||
callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
|
||||
else:
|
||||
callbacks = [StdOutCallbackHandler()]
|
||||
|
||||
llm = Ollama(
|
||||
model=config.model,
|
||||
system=config.system_prompt,
|
||||
temperature=config.temperature,
|
||||
top_p=config.top_p,
|
||||
callback_manager=CallbackManager(callbacks),
|
||||
base_url=config.base_url,
|
||||
)
|
||||
|
||||
return llm.invoke(prompt)
|
||||
@@ -0,0 +1,107 @@
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Callable, Dict, Optional, Type, Union
|
||||
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
from langchain.schema import BaseMessage, HumanMessage, SystemMessage
|
||||
from langchain_core.tools import BaseTool
|
||||
from langchain_openai import ChatOpenAI
|
||||
from pydantic import BaseModel
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class OpenAILlm(BaseLlm):
|
||||
def __init__(
|
||||
self,
|
||||
config: Optional[BaseLlmConfig] = None,
|
||||
tools: Optional[Union[Dict[str, Any], Type[BaseModel], Callable[..., Any], BaseTool]] = None,
|
||||
):
|
||||
self.tools = tools
|
||||
super().__init__(config=config)
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "openai/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["completion_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["prompt_tokens"],
|
||||
"completion_tokens": token_info["completion_tokens"],
|
||||
"total_tokens": token_info["prompt_tokens"] + token_info["completion_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
def _get_answer(self, prompt: str, config: BaseLlmConfig) -> str:
|
||||
messages = []
|
||||
if config.system_prompt:
|
||||
messages.append(SystemMessage(content=config.system_prompt))
|
||||
messages.append(HumanMessage(content=prompt))
|
||||
kwargs = {
|
||||
"model": config.model or "gpt-3.5-turbo",
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"model_kwargs": config.model_kwargs or {},
|
||||
}
|
||||
api_key = config.api_key or os.environ["OPENAI_API_KEY"]
|
||||
base_url = config.base_url or os.environ.get("OPENAI_API_BASE", None)
|
||||
if config.top_p:
|
||||
kwargs["model_kwargs"]["top_p"] = config.top_p
|
||||
if config.default_headers:
|
||||
kwargs["default_headers"] = config.default_headers
|
||||
if config.stream:
|
||||
callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
|
||||
chat = ChatOpenAI(
|
||||
**kwargs,
|
||||
streaming=config.stream,
|
||||
callbacks=callbacks,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
http_client=config.http_client,
|
||||
http_async_client=config.http_async_client,
|
||||
)
|
||||
else:
|
||||
chat = ChatOpenAI(
|
||||
**kwargs,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
http_client=config.http_client,
|
||||
http_async_client=config.http_async_client,
|
||||
)
|
||||
if self.tools:
|
||||
return self._query_function_call(chat, self.tools, messages)
|
||||
|
||||
chat_response = chat.invoke(messages)
|
||||
if self.config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
|
||||
def _query_function_call(
|
||||
self,
|
||||
chat: ChatOpenAI,
|
||||
tools: Optional[Union[Dict[str, Any], Type[BaseModel], Callable[..., Any], BaseTool]],
|
||||
messages: list[BaseMessage],
|
||||
) -> str:
|
||||
from langchain.output_parsers.openai_tools import JsonOutputToolsParser
|
||||
from langchain_core.utils.function_calling import convert_to_openai_tool
|
||||
|
||||
openai_tools = [convert_to_openai_tool(tools)]
|
||||
chat = chat.bind(tools=openai_tools).pipe(JsonOutputToolsParser())
|
||||
try:
|
||||
return json.dumps(chat.invoke(messages)[0])
|
||||
except IndexError:
|
||||
return "Input could not be mapped to the function!"
|
||||
@@ -0,0 +1,71 @@
|
||||
import importlib
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
try:
|
||||
from langchain_together import ChatTogether
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Please install the langchain_together package by running `pip install langchain_together==0.1.3`."
|
||||
)
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class TogetherLlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
try:
|
||||
importlib.import_module("together")
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for Together are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[together]"`'
|
||||
) from None
|
||||
|
||||
super().__init__(config=config)
|
||||
if not self.config.api_key and "TOGETHER_API_KEY" not in os.environ:
|
||||
raise ValueError("Please set the TOGETHER_API_KEY environment variable or pass it in the config.")
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.system_prompt:
|
||||
raise ValueError("TogetherLlm does not support `system_prompt`")
|
||||
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "together/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_tokens"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info["completion_tokens"]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["prompt_tokens"],
|
||||
"completion_tokens": token_info["completion_tokens"],
|
||||
"total_tokens": token_info["prompt_tokens"] + token_info["completion_tokens"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
api_key = config.api_key or os.environ["TOGETHER_API_KEY"]
|
||||
kwargs = {
|
||||
"model_name": config.model or "mixtral-8x7b-32768",
|
||||
"temperature": config.temperature,
|
||||
"max_tokens": config.max_tokens,
|
||||
"together_api_key": api_key,
|
||||
}
|
||||
|
||||
chat = ChatTogether(**kwargs)
|
||||
chat_response = chat.invoke(prompt)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["token_usage"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,68 @@
|
||||
import importlib
|
||||
import logging
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
from langchain_google_vertexai import ChatVertexAI
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class VertexAILlm(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
try:
|
||||
importlib.import_module("vertexai")
|
||||
except ModuleNotFoundError:
|
||||
raise ModuleNotFoundError(
|
||||
"The required dependencies for VertexAI are not installed."
|
||||
'Please install with `pip install --upgrade "embedchain[vertexai]"`'
|
||||
) from None
|
||||
super().__init__(config=config)
|
||||
|
||||
def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
|
||||
if self.config.token_usage:
|
||||
response, token_info = self._get_answer(prompt, self.config)
|
||||
model_name = "vertexai/" + self.config.model
|
||||
if model_name not in self.config.model_pricing_map:
|
||||
raise ValueError(
|
||||
f"Model {model_name} not found in `model_prices_and_context_window.json`. \
|
||||
You can disable token usage by setting `token_usage` to False."
|
||||
)
|
||||
total_cost = (
|
||||
self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_token_count"]
|
||||
) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info[
|
||||
"candidates_token_count"
|
||||
]
|
||||
response_token_info = {
|
||||
"prompt_tokens": token_info["prompt_token_count"],
|
||||
"completion_tokens": token_info["candidates_token_count"],
|
||||
"total_tokens": token_info["prompt_token_count"] + token_info["candidates_token_count"],
|
||||
"total_cost": round(total_cost, 10),
|
||||
"cost_currency": "USD",
|
||||
}
|
||||
return response, response_token_info
|
||||
return self._get_answer(prompt, self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
|
||||
if config.top_p and config.top_p != 1:
|
||||
logger.warning("Config option `top_p` is not supported by this model.")
|
||||
|
||||
if config.stream:
|
||||
callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
|
||||
llm = ChatVertexAI(
|
||||
temperature=config.temperature, model=config.model, callbacks=callbacks, streaming=config.stream
|
||||
)
|
||||
else:
|
||||
llm = ChatVertexAI(temperature=config.temperature, model=config.model)
|
||||
|
||||
messages = VertexAILlm._get_messages(prompt)
|
||||
chat_response = llm.invoke(messages)
|
||||
if config.token_usage:
|
||||
return chat_response.content, chat_response.response_metadata["usage_metadata"]
|
||||
return chat_response.content
|
||||
@@ -0,0 +1,40 @@
|
||||
from typing import Iterable, Optional, Union
|
||||
|
||||
from langchain.callbacks.manager import CallbackManager
|
||||
from langchain.callbacks.stdout import StdOutCallbackHandler
|
||||
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||||
from langchain_community.llms import VLLM as BaseVLLM
|
||||
|
||||
from embedchain.config import BaseLlmConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
from embedchain.llm.base import BaseLlm
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class VLLM(BaseLlm):
|
||||
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
||||
super().__init__(config=config)
|
||||
if self.config.model is None:
|
||||
self.config.model = "mosaicml/mpt-7b"
|
||||
|
||||
def get_llm_model_answer(self, prompt):
|
||||
return self._get_answer(prompt=prompt, config=self.config)
|
||||
|
||||
@staticmethod
|
||||
def _get_answer(prompt: str, config: BaseLlmConfig) -> Union[str, Iterable]:
|
||||
callback_manager = [StreamingStdOutCallbackHandler()] if config.stream else [StdOutCallbackHandler()]
|
||||
|
||||
# Prepare the arguments for BaseVLLM
|
||||
llm_args = {
|
||||
"model": config.model,
|
||||
"temperature": config.temperature,
|
||||
"top_p": config.top_p,
|
||||
"callback_manager": CallbackManager(callback_manager),
|
||||
}
|
||||
|
||||
# Add model_kwargs if they are not None
|
||||
if config.model_kwargs is not None:
|
||||
llm_args.update(config.model_kwargs)
|
||||
|
||||
llm = BaseVLLM(**llm_args)
|
||||
return llm.invoke(prompt)
|
||||
Reference in New Issue
Block a user