klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Show details for query tokens (#1392)

Browse Source
This commit is contained in:
Dev Khant
2024-07-05 00:10:56 +05:30
committed by GitHub
parent ea09b5f7f0
commit 4880557d51
25 changed files with 1825 additions and 517 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
embedchain/llm/anthropic.py
View File

@@ -1,6 +1,6 @@
import logging
import os
from typing import Optional
from typing import Any, Optional
try:
from langchain_anthropic import ChatAnthropic
@@ -21,8 +21,27 @@ class AnthropicLlm(BaseLlm):
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):
return AnthropicLlm._get_answer(prompt=prompt, config=self.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:
@@ -34,4 +53,7 @@ class AnthropicLlm(BaseLlm):
messages = BaseLlm._get_messages(prompt, system_prompt=config.system_prompt)
return chat(messages).content
chat_response = chat.invoke(messages)
if config.token_usage:
return chat_response.content, chat_response.response_metadata["token_usage"]
return chat_response.content

19
embedchain/llm/base.py
View File

@@ -164,7 +164,7 @@ class BaseLlm(JSONSerializable):
return search.run(input_query)
@staticmethod
def _stream_response(answer: Any) -> Generator[Any, Any, None]:
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
@@ -177,6 +177,8 @@ class BaseLlm(JSONSerializable):
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):
"""
@@ -219,11 +221,18 @@ class BaseLlm(JSONSerializable):
if dry_run:
return prompt
answer = self.get_answer_from_llm(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:
@@ -276,13 +285,13 @@ class BaseLlm(JSONSerializable):
if dry_run:
return prompt
answer = self.get_answer_from_llm(prompt)
answer, token_info = self.get_answer_from_llm(prompt)
if isinstance(answer, str):
logger.info(f"Answer: {answer}")
return answer
return answer, token_info
else:
# this is a streamed response and needs to be handled differently.
return self._stream_response(answer)
return self._stream_response(answer, token_info)
finally:
if config:
# Restore previous config

51
embedchain/llm/cohere.py
View File

@@ -1,8 +1,8 @@
import importlib
import os
from typing import Optional
from typing import Any, Optional
from langchain_community.llms.cohere import Cohere
from langchain_cohere import ChatCohere
from embedchain.config import BaseLlmConfig
from embedchain.helpers.json_serializable import register_deserializable
@@ -17,27 +17,50 @@ class CohereLlm(BaseLlm):
except ModuleNotFoundError:
raise ModuleNotFoundError(
"The required dependencies for Cohere are not installed."
'Please install with `pip install --upgrade "embedchain[cohere]"`'
"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):
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`")
return CohereLlm._get_answer(prompt=prompt, config=self.config)
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.getenv("COHERE_API_KEY")
llm = Cohere(
cohere_api_key=api_key,
model=config.model,
max_tokens=config.max_tokens,
temperature=config.temperature,
p=config.top_p,
)
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,
}
return llm.invoke(prompt)
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

32
embedchain/llm/groq.py
View File

@@ -1,5 +1,5 @@
import os
from typing import Optional
from typing import Any, Optional
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.schema import HumanMessage, SystemMessage
@@ -22,9 +22,27 @@ class GroqLlm(BaseLlm):
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) -> str:
response = self._get_answer(prompt, self.config)
return response
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 = []
@@ -42,4 +60,8 @@ class GroqLlm(BaseLlm):
chat = ChatGroq(**kwargs, streaming=config.stream, callbacks=callbacks, api_key=api_key)
else:
chat = ChatGroq(**kwargs)
return chat.invoke(messages).content
chat_response = chat.invoke(prompt)
if self.config.token_usage:
return chat_response.content, chat_response.response_metadata["token_usage"]
return chat_response.content

32
embedchain/llm/mistralai.py
View File

@@ -1,5 +1,5 @@
import os
from typing import Optional
from typing import Any, Optional
from embedchain.config import BaseLlmConfig
from embedchain.helpers.json_serializable import register_deserializable
@@ -13,8 +13,27 @@ class MistralAILlm(BaseLlm):
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):
return MistralAILlm._get_answer(prompt=prompt, config=self.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):
@@ -47,6 +66,7 @@ class MistralAILlm(BaseLlm):
answer += chunk.content
return answer
else:
response = client.invoke(**kwargs, input=messages)
answer = response.content
return answer
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

30
embedchain/llm/nvidia.py
View File

@@ -1,6 +1,6 @@
import os
from collections.abc import Iterable
from typing import Optional, Union
from typing import Any, Optional, Union
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.stdout import StdOutCallbackHandler
@@ -25,8 +25,27 @@ class NvidiaLlm(BaseLlm):
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):
return self._get_answer(prompt=prompt, config=self.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]:
@@ -43,4 +62,7 @@ class NvidiaLlm(BaseLlm):
if labels:
params["labels"] = labels
llm = ChatNVIDIA(**params, callback_manager=CallbackManager(callback_manager))
return llm.invoke(prompt).content if labels is None else llm.invoke(prompt, labels=labels).content
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

30
embedchain/llm/openai.py
View File

@@ -23,9 +23,28 @@ class OpenAILlm(BaseLlm):
self.tools = tools
super().__init__(config=config)
def get_llm_model_answer(self, prompt) -> str:
response = self._get_answer(prompt, self.config)
return response
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 = []
@@ -66,7 +85,10 @@ class OpenAILlm(BaseLlm):
if self.tools:
return self._query_function_call(chat, self.tools, messages)
return chat.invoke(messages).content
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,

54
embedchain/llm/together.py
View File

@@ -1,8 +1,13 @@
import importlib
import os
from typing import Optional
from typing import Any, Optional
from langchain_community.llms import Together
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
@@ -24,20 +29,43 @@ class TogetherLlm(BaseLlm):
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):
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`")
return TogetherLlm._get_answer(prompt=prompt, config=self.config)
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.getenv("TOGETHER_API_KEY")
llm = Together(
together_api_key=api_key,
model=config.model,
max_tokens=config.max_tokens,
temperature=config.temperature,
top_p=config.top_p,
)
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,
}
return llm.invoke(prompt)
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

35
embedchain/llm/vertex_ai.py
View File

@@ -1,6 +1,6 @@
import importlib
import logging
from typing import Optional
from typing import Any, Optional
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain_google_vertexai import ChatVertexAI
@@ -24,16 +24,35 @@ class VertexAILlm(BaseLlm):
) from None
super().__init__(config=config)
def get_llm_model_answer(self, prompt):
return VertexAILlm._get_answer(prompt=prompt, config=self.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.")
messages = BaseLlm._get_messages(prompt, system_prompt=config.system_prompt)
if config.stream:
callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
llm = ChatVertexAI(
@@ -42,4 +61,8 @@ class VertexAILlm(BaseLlm):
else:
llm = ChatVertexAI(temperature=config.temperature, model=config.model)
return llm.invoke(messages).content
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