Compare commits

..
11 Commits
50 changed files with 2107 additions and 724 deletions
+4 -1
View File
@@ -76,7 +76,7 @@ docs/_build/
target/
# Jupyter Notebook
.ipynb_checkpoints
*.yaml
# IPython
profile_default/
@@ -171,3 +171,6 @@ db
.idea/
.DS_Store
notebooks/*.yaml
.ipynb_checkpoints/
+2 -2
View File
@@ -8,7 +8,7 @@
[![Open in Colab](https://camo.githubusercontent.com/84f0493939e0c4de4e6dbe113251b4bfb5353e57134ffd9fcab6b8714514d4d1/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/drive/138lMWhENGeEu7Q1-6lNbNTHGLZXBBz_B?usp=sharing)
[![codecov](https://codecov.io/gh/embedchain/embedchain/graph/badge.svg?token=EMRRHZXW1Q)](https://codecov.io/gh/embedchain/embedchain)
Embedchain is a framework to easily create LLM powered bots over any dataset. If you want a javascript version, check out [embedchain-js](https://github.com/embedchain/embedchain/tree/main/embedchain-js)
Embedchain is a Data Platform for LLMs - load, index, retrieve, and sync any unstructured data. Using embedchain, you can easily create LLM powered apps over any data. If you want a javascript version, check out [embedchain-js](https://github.com/embedchain/embedchain/tree/main/embedchain-js)
## Community
@@ -94,7 +94,7 @@ If you utilize this repository, please consider citing it with:
```
@misc{embedchain,
author = {Taranjeet Singh, Deshraj Yadav},
title = {Embedchain: Framework to easily create LLM powered bots over any dataset},
title = {Embedchain: Data platform for LLMs - load, index, retrieve, and sync any unstructured data},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
+19
View File
@@ -0,0 +1,19 @@
app:
config:
id: azure-openai-app
llm:
provider: azure_openai
model: gpt-35-turbo
config:
deployment_name: your_llm_deployment_name
temperature: 0.5
max_tokens: 1000
top_p: 1
stream: false
embedder:
provider: azure_openai
config:
model: text-embedding-ada-002
deployment_name: you_embedding_model_deployment_name
+1 -1
View File
@@ -25,7 +25,7 @@ llm:
vectordb:
provider: chroma
config:
collection_name: 'full-stack-app'
collection_name: 'my-collection-name'
dir: db
allow_reset: true
+6
View File
@@ -0,0 +1,6 @@
vectordb:
provider: pinecone
config:
metric: cosine
vector_dimension: 1536
collection_name: my-pinecone-index
+4
View File
@@ -0,0 +1,4 @@
vectordb:
provider: weaviate
config:
collection_name: my_weaviate_index
-66
View File
@@ -1,66 +0,0 @@
---
title: '🔍 Query configurations'
---
## AppConfig
| option | description | type | default |
|-----------|-----------------------|---------------------------------|------------------------|
| log_level | log level | string | WARNING |
| embedding_fn| embedding function | chromadb.utils.embedding_functions | \{text-embedding-ada-002\} |
| db | vector database (experimental) | BaseVectorDB | ChromaDB |
| collection_name | initial collection name for the database | string | embedchain_store |
| collect_metrics | collect anonymous telemetry data to improve embedchain | boolean | true |
## AddConfig
|option|description|type|default|
|---|---|---|---|
|chunker|chunker config|ChunkerConfig|Default values for chunker depends on the `data_type`. Please refer [ChunkerConfig](#chunker-config)|
|loader|loader config|LoaderConfig|None|
Yes, you are passing `ChunkerConfig` to `AddConfig`, like so:
```python
chunker_config = ChunkerConfig(chunk_size=100)
add_config = AddConfig(chunker=chunker_config)
app.add("lorem ipsum", config=add_config)
```
### ChunkerConfig
|option|description|type|default|
|---|---|---|---|
|chunk_size|Maximum size of chunks to return|int|Default value for various `data_type` mentioned below|
|chunk_overlap|Overlap in characters between chunks|int|Default value for various `data_type` mentioned below|
|length_function|Function that measures the length of given chunks|typing.Callable|Default value for various `data_type` mentioned below|
Default values of chunker config parameters for different `data_type`:
|data_type|chunk_size|chunk_overlap|length_function|
|---|---|---|---|
|docx|1000|0|len|
|text|300|0|len|
|qna_pair|300|0|len|
|web_page|500|0|len|
|pdf_file|1000|0|len|
|youtube_video|2000|0|len|
|docs_site|500|50|len|
|notion|300|0|len|
## BaseLlmConfig
|option|description|type|default|
|---|---|---|---|
|number_documents|Absolute number of documents to pull from the database as context.|int|1
|template|custom template for prompt. If history is used with query, $history has to be included as well.|Template|Template("Use the following pieces of context to answer the query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. \$context Query: \$query Helpful Answer:")|
|model|name of the model used.|string|depends on app type|
|temperature|Controls the randomness of the model's output. Higher values (closer to 1) make output more random, lower values make it more deterministic.|float|0|
|max_tokens|Controls how many tokens are used. Exact implementation (whether it counts prompt and/or response) depends on the model.|int|1000|
|top_p|Controls the diversity of words. Higher values (closer to 1) make word selection more diverse, lower values make words less diverse.|float|1|
|history|include conversation history from your client or database.|any (recommendation: list[str])|None|
|stream|control if response is streamed back to the user.|bool|False|
|deployment_name|t.b.a.|str|None|
|system_prompt|System prompt string. Unused if none.|str|None|
|where|filter for context search.|dict|None|
-40
View File
@@ -1,40 +0,0 @@
---
title: '🧪 Testing'
---
## Methods for testing
### Dry Run
Before you consume valueable tokens, you should make sure that data chunks are properly created and the embedding you have done works and that it's receiving the correct document from the database.
- For `query` or `chat` method, you can add this to your script:
```python
print(naval_chat_bot.query('Can you tell me who Naval Ravikant is?', dry_run=True))
'''
Use the following pieces of context to answer the query at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
Q: Who is Naval Ravikant?
A: Naval Ravikant is an Indian-American entrepreneur and investor.
Query: Can you tell me who Naval Ravikant is?
Helpful Answer:
'''
```
_The embedding is confirmed to work as expected. It returns the right document, even if the question is asked slightly different. No prompt tokens have been consumed._
The dry run will still consume tokens to embed your query, but it is only **~1/15 of the prompt.**
- For `add` method, you can add this to your script:
```python
print(naval_chat_bot.add('https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Almanack-of-Naval-Ravikant_Final.pdf', dry_run=True))
'''
{'chunks': ['THE ALMANACK OF NAVAL RAVIKANT', 'GETTING RICH IS NOT JUST ABOUT LUCK;', 'HAPPINESS IS NOT JUST A TRAIT WE ARE'], 'metadata': [{'source': 'C:\\Users\\Dev\\AppData\\Local\\Temp\\tmp3g5mjoiz\\tmp.pdf', 'page': 0, 'url': 'https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Almanack-of-Naval-Ravikant_Final.pdf', 'data_type': 'pdf_file'}, {'source': 'C:\\Users\\Dev\\AppData\\Local\\Temp\\tmp3g5mjoiz\\tmp.pdf', 'page': 2, 'url': 'https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Almanack-of-Naval-Ravikant_Final.pdf', 'data_type': 'pdf_file'}, {'source': 'C:\\Users\\Dev\\AppData\\Local\\Temp\\tmp3g5mjoiz\\tmp.pdf', 'page': 2, 'url': 'https://navalmanack.s3.amazonaws.com/Eric-Jorgenson_The-Almanack-of-Naval-Ravikant_Final.pdf', 'data_type': 'pdf_file'}], 'count': 7358, 'type': <DataType.PDF_FILE: 'pdf_file'>}
# less items to show for readability
'''
```
-118
View File
@@ -1,118 +0,0 @@
---
title: '💾 Vector Database'
---
We support `Chroma`, `Elasticsearch` and `OpenSearch` as vector databases.
`Chroma` is used as a default database.
## Elasticsearch
### Minimal Example
In order to use `Elasticsearch` as vector database we need to use App type `CustomApp`.
1. Set the environment variables in a `.env` file.
```
OPENAI_API_KEY=sk-SECRETKEY
ELASTICSEARCH_API_KEY=SECRETKEY==
ELASTICSEARCH_URL=https://secret-domain.europe-west3.gcp.cloud.es.io:443
```
Please note that the key needs certain privileges. For testing you can just toggle off `restrict privileges` under `/app/management/security/api_keys/` in your web interface.
2. Load the app
```python
from embedchain import CustomApp
from embedchain.embedder.openai import OpenAIEmbedder
from embedchain.llm.openai import OpenAILlm
from embedchain.vectordb.elasticsearch import ElasticsearchDB
es_app = CustomApp(
llm=OpenAILlm(),
embedder=OpenAIEmbedder(),
db=ElasticsearchDB(),
)
```
### More custom settings
You can get a URL for elasticsearch in the cloud, or run it locally.
The following example shows you how to configure embedchain to work with a locally running elasticsearch.
Instead of using an API key, we use http login credentials. The localhost url can be defined in .env or in the config.
```python
import os
from embedchain import CustomApp
from embedchain.config import CustomAppConfig, ElasticsearchDBConfig
from embedchain.embedder.openai import OpenAIEmbedder
from embedchain.llm.openai import OpenAILlm
from embedchain.vectordb.elasticsearch import ElasticsearchDB
es_config = ElasticsearchDBConfig(
# elasticsearch url or list of nodes url with different hosts and ports.
es_url='https://localhost:9200',
# pass named parameters supported by Python Elasticsearch client
http_auth=("elastic", "secret"),
ca_certs="~/binaries/elasticsearch-8.7.0/config/certs/http_ca.crt" # your cert path
# verify_certs=False # Alternative, if you aren't using certs
) # pass named parameters supported by elasticsearch-py
es_app = CustomApp(
config=CustomAppConfig(log_level="INFO"),
llm=OpenAILlm(),
embedder=OpenAIEmbedder(),
db=ElasticsearchDB(config=es_config),
)
```
3. This should log your connection details to the console.
4. Alternatively to a URL, you `ElasticsearchDBConfig` accepts `es_url` as a list of nodes url with different hosts and ports.
5. Additionally we can pass named parameters supported by Python Elasticsearch client.
## OpenSearch 🔍
To use OpenSearch as a vector database with a CustomApp, follow these simple steps:
1. Set the `OPENAI_API_KEY` environment variable:
```
OPENAI_API_KEY=sk-xxxx
```
2. Define the OpenSearch configuration in your Python code:
```python
from embedchain import CustomApp
from embedchain.config import OpenSearchDBConfig
from embedchain.embedder.openai import OpenAIEmbedder
from embedchain.llm.openai import OpenAILlm
from embedchain.vectordb.opensearch import OpenSearchDB
opensearch_url = "https://localhost:9200"
http_auth = ("username", "password")
db_config = OpenSearchDBConfig(
opensearch_url=opensearch_url,
http_auth=http_auth,
collection_name="embedchain-app",
use_ssl=True,
timeout=30,
)
db = OpenSearchDB(config=db_config)
```
2. Instantiate the app and add data:
```python
app = CustomApp(llm=OpenAILlm(), embedder=OpenAIEmbedder(), db=db)
app.add("https://en.wikipedia.org/wiki/Elon_Musk")
app.add("https://www.forbes.com/profile/elon-musk")
app.add("https://www.britannica.com/biography/Elon-Musk")
```
3. You're all set! Start querying using the following command:
```python
app.query("What is the net worth of Elon Musk?")
```
+20 -10
View File
@@ -4,15 +4,25 @@ title: 🤝 Connect with Us
We believe in building a vibrant and supportive community around embedchain. There are various channels through which you can connect with us, stay updated, and contribute to the ongoing discussions:
* Slack: Our Slack workspace provides a platform for more structured discussions and channels dedicated to different topics. Feel free to jump in and start contributing. [Join Slack](https://join.slack.com/t/embedchain/shared_invite/zt-22uwz3c46-Zg7cIh5rOBteT_xe1jwLDw).
* Discord: Join our Discord server to engage in real-time conversations with the community members and the project maintainers. It’s a great place to seek help and discuss anything related to the project. [Join Discord](https://discord.gg/CUU9FPhRNt).
* Twitter: Follow us on Twitter for the latest news, announcements, and highlights from our community. It’s also a quick way to reach out to us. [Follow @embedchain](https://twitter.com/embedchain).
* LinkedIn: Connect with us on LinkedIn to stay updated on official announcements, job openings, and professional networking opportunities within our community. [Follow Our Page](https://www.linkedin.com/company/embedchain/).
* Newsletter: Subscribe to our newsletter for a curated list of project updates, community contributions, and upcoming events. It’s a compact way to stay in the loop with what’s happening in our community. [Subscribe Now](https://embedchain.substack.com/).
<CardGroup cols={3}>
<Card title="Twitter" icon="twitter" href="https://twitter.com/embedchain">
Follow us on Twitter
</Card>
<Card title="Slack" icon="slack" href="https://join.slack.com/t/embedchain/shared_invite/zt-22uwz3c46-Zg7cIh5rOBteT_xe1jwLDw" color="#4A154B">
Join our slack community
</Card>
<Card title="Discord" icon="discord" href="https://discord.gg/6PzXDgEjG5" color="#7289DA">
Join our discord community
</Card>
<Card title="LinkedIn" icon="linkedin" href="https://www.linkedin.com/company/embedchain/">
Connect with us on LinkedIn
</Card>
<Card title="Schedule a call" icon="calendar" href="https://cal.com/taranjeetio/ec">
Schedule a call with Embedchain founder
</Card>
<Card title="Newsletter" icon="message" href="https://embedchain.substack.com/">
Subscribe to our newsletter
</Card>
</CardGroup>
We look forward to connecting with you and seeing how we can create amazing things together!
+52 -12
View File
@@ -8,6 +8,7 @@ Embedchain supports several embedding models from the following providers:
<CardGroup cols={4}>
<Card title="OpenAI" href="#openai"></Card>
<Card title="Azure OpenAI" href="#azure-openai"></Card>
<Card title="GPT4All" href="#gpt4all"></Card>
<Card title="Hugging Face" href="#hugging-face"></Card>
<Card title="Vertex AI" href="#vertex-ai"></Card>
@@ -27,14 +28,14 @@ from embedchain import App
os.environ['OPENAI_API_KEY'] = 'xxx'
# load embedding model configuration from openai.yaml file
app = App.from_config(yaml_path="openai.yaml")
# load embedding model configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
app.add("https://en.wikipedia.org/wiki/OpenAI")
app.query("What is OpenAI?")
```
```yaml openai.yaml
```yaml config.yaml
embedder:
provider: openai
config:
@@ -43,6 +44,45 @@ embedder:
</CodeGroup>
## Azure OpenAI
To use Azure OpenAI embedding model, you have to set some of the azure openai related environment variables as given in the code block below:
<CodeGroup>
```python main.py
import os
from embedchain import App
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_BASE"] = "https://xxx.openai.azure.com/"
os.environ["OPENAI_API_KEY"] = "xxx"
os.environ["OPENAI_API_VERSION"] = "xxx"
app = App.from_config(yaml_path="config.yaml")
```
```yaml config.yaml
llm:
provider: azure_openai
model: gpt-35-turbo
config:
deployment_name: your_llm_deployment_name
temperature: 0.5
max_tokens: 1000
top_p: 1
stream: false
embedder:
provider: azure_openai
config:
model: text-embedding-ada-002
deployment_name: you_embedding_model_deployment_name
```
</CodeGroup>
You can find the list of models and deployment name on the [Azure OpenAI Platform](https://oai.azure.com/portal).
## GPT4ALL
GPT4All supports generating high quality embeddings of arbitrary length documents of text using a CPU optimized contrastively trained Sentence Transformer.
@@ -52,11 +92,11 @@ GPT4All supports generating high quality embeddings of arbitrary length document
```python main.py
from embedchain import App
# load embedding model configuration from gpt4all.yaml file
app = App.from_config(yaml_path="gpt4all.yaml")
# load embedding model configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml gpt4all.yaml
```yaml config.yaml
llm:
provider: gpt4all
model: 'orca-mini-3b.ggmlv3.q4_0.bin'
@@ -83,11 +123,11 @@ Hugging Face supports generating embeddings of arbitrary length documents of tex
```python main.py
from embedchain import App
# load embedding model configuration from huggingface.yaml file
app = App.from_config(yaml_path="huggingface.yaml")
# load embedding model configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml huggingface.yaml
```yaml config.yaml
llm:
provider: huggingface
model: 'google/flan-t5-xxl'
@@ -114,11 +154,11 @@ Embedchain supports Google's VertexAI embeddings model through a simple interfac
```python main.py
from embedchain import App
# load embedding model configuration from vertexai.yaml file
app = App.from_config(yaml_path="vertexai.yaml")
# load embedding model configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml vertexai.yaml
```yaml config.yaml
llm:
provider: vertexai
model: 'chat-bison'
+80 -33
View File
@@ -35,7 +35,7 @@ app.add("https://en.wikipedia.org/wiki/OpenAI")
app.query("What is OpenAI?")
```
If you are looking to configure the different parameters of the LLM, you can do so by loading the app using a [yaml config](https://github.com/embedchain/embedchain/blob/main/embedchain/yaml/chroma.yaml) file.
If you are looking to configure the different parameters of the LLM, you can do so by loading the app using a [yaml config](https://github.com/embedchain/embedchain/blob/main/configs/chroma.yaml) file.
<CodeGroup>
@@ -45,11 +45,11 @@ from embedchain import App
os.environ['OPENAI_API_KEY'] = 'xxx'
# load llm configuration from openai.yaml file
app = App.from_config(yaml_path="openai.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml openai.yaml
```yaml config.yaml
llm:
provider: openai
model: 'gpt-3.5-turbo'
@@ -65,7 +65,42 @@ llm:
## Azure OpenAI
_Coming soon_
To use Azure OpenAI model, you have to set some of the azure openai related environment variables as given in the code block below:
<CodeGroup>
```python main.py
import os
from embedchain import App
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_BASE"] = "https://xxx.openai.azure.com/"
os.environ["OPENAI_API_KEY"] = "xxx"
os.environ["OPENAI_API_VERSION"] = "xxx"
app = App.from_config(yaml_path="config.yaml")
```
```yaml config.yaml
llm:
provider: azure_openai
model: gpt-35-turbo
config:
deployment_name: your_llm_deployment_name
temperature: 0.5
max_tokens: 1000
top_p: 1
stream: false
embedder:
provider: azure_openai
config:
model: text-embedding-ada-002
deployment_name: you_embedding_model_deployment_name
```
</CodeGroup>
You can find the list of models and deployment name on the [Azure OpenAI Platform](https://oai.azure.com/portal).
## Anthropic
@@ -79,11 +114,11 @@ from embedchain import App
os.environ["ANTHROPIC_API_KEY"] = "xxx"
# load llm configuration from anthropic.yaml file
app = App.from_config(yaml_path="anthropic.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml anthropic.yaml
```yaml config.yaml
llm:
provider: anthropic
model: 'claude-instant-1'
@@ -96,15 +131,14 @@ llm:
</CodeGroup>
<br />
<Tip>
You may also have to set the `OPENAI_API_KEY` if you use the OpenAI's embedding model.
</Tip>
## Cohere
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[cohere]'
```
Set the `COHERE_API_KEY` as environment variable which you can find on their [Account settings page](https://dashboard.cohere.com/api-keys).
Once you have the API key, you are all set to use it with Embedchain.
@@ -117,11 +151,11 @@ from embedchain import App
os.environ["COHERE_API_KEY"] = "xxx"
# load llm configuration from cohere.yaml file
app = App.from_config(yaml_path="cohere.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml cohere.yaml
```yaml config.yaml
llm:
provider: cohere
model: large
@@ -135,6 +169,12 @@ llm:
## GPT4ALL
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[opensource]'
```
GPT4all is a free-to-use, locally running, privacy-aware chatbot. No GPU or internet required. You can use this with Embedchain using the following code:
<CodeGroup>
@@ -142,11 +182,11 @@ GPT4all is a free-to-use, locally running, privacy-aware chatbot. No GPU or inte
```python main.py
from embedchain import App
# load llm configuration from gpt4all.yaml file
app = App.from_config(yaml_path="gpt4all.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml gpt4all.yaml
```yaml config.yaml
llm:
provider: gpt4all
model: 'orca-mini-3b.ggmlv3.q4_0.bin'
@@ -177,11 +217,11 @@ import os
from embedchain import App
os.environ["JINACHAT_API_KEY"] = "xxx"
# load llm configuration from jina.yaml file
app = App.from_config(yaml_path="jina.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml jina.yaml
```yaml config.yaml
llm:
provider: jina
config:
@@ -195,6 +235,13 @@ llm:
## Hugging Face
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[huggingface_hub]'
```
First, set `HUGGINGFACE_ACCESS_TOKEN` in environment variable which you can obtain from [their platform](https://huggingface.co/settings/tokens).
Once you have the token, load the app using the config yaml file:
@@ -207,11 +254,11 @@ from embedchain import App
os.environ["HUGGINGFACE_ACCESS_TOKEN"] = "xxx"
# load llm configuration from huggingface.yaml file
app = App.from_config(yaml_path="huggingface.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml huggingface.yaml
```yaml config.yaml
llm:
provider: huggingface
model: 'google/flan-t5-xxl'
@@ -237,11 +284,11 @@ from embedchain import App
os.environ["REPLICATE_API_TOKEN"] = "xxx"
# load llm configuration from llama2.yaml file
app = App.from_config(yaml_path="llama2.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml llama2.yaml
```yaml config.yaml
llm:
provider: llama2
model: 'a16z-infra/llama13b-v2-chat:df7690f1994d94e96ad9d568eac121aecf50684a0b0963b25a41cc40061269e5'
@@ -262,11 +309,11 @@ Setup Google Cloud Platform application credentials by following the instruction
```python main.py
from embedchain import App
# load llm configuration from vertexai.yaml file
app = App.from_config(yaml_path="vertexai.yaml")
# load llm configuration from config.yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml vertexai.yaml
```yaml config.yaml
llm:
provider: vertexai
model: 'chat-bison'
+76 -11
View File
@@ -25,10 +25,10 @@ Utilizing a vector database alongside Embedchain is a seamless process. All you
from embedchain import App
# load chroma configuration from yaml file
app = App.from_config(yaml_path="chroma-config-1.yaml")
app = App.from_config(yaml_path="config1.yaml")
```
```yaml chroma-config-1.yaml
```yaml config1.yaml
vectordb:
provider: chroma
config:
@@ -37,7 +37,7 @@ vectordb:
allow_reset: true
```
```yaml chroma-config-2.yaml
```yaml config2.yaml
vectordb:
provider: chroma
config:
@@ -52,16 +52,22 @@ vectordb:
## Elasticsearch
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[elasticsearch]'
```
<CodeGroup>
```python main.py
from embedchain import App
# load elasticsearch configuration from yaml file
app = App.from_config(yaml_path="elasticsearch.yaml")
app = App.from_config(yaml_path="config.yaml")
```
```yaml elasticsearch.yaml
```yaml config.yaml
vectordb:
provider: elasticsearch
config:
@@ -74,16 +80,22 @@ vectordb:
## OpenSearch
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[opensearch]'
```
<CodeGroup>
```python main.py
from embedchain import App
# load opensearch configuration from yaml file
app = App.from_config(yaml_path="opensearch.yaml")
app = App.from_config(yaml_path="config.yaml")
```
```yaml opensearch.yaml
```yaml config.yaml
vectordb:
provider: opensearch
config:
@@ -101,16 +113,28 @@ vectordb:
## Zilliz
Install related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[milvus]'
```
Set the Zilliz environment variables `ZILLIZ_CLOUD_URI` and `ZILLIZ_CLOUD_TOKEN` which you can find it on their [cloud platform](https://cloud.zilliz.com/).
<CodeGroup>
```python main.py
import os
from embedchain import App
os.environ['ZILLIZ_CLOUD_URI'] = 'https://xxx.zillizcloud.com'
os.environ['ZILLIZ_CLOUD_TOKEN'] = 'xxx'
# load zilliz configuration from yaml file
app = App.from_config(yaml_path="zilliz.yaml")
app = App.from_config(yaml_path="config.yaml")
```
```yaml zilliz.yaml
```yaml config.yaml
vectordb:
provider: zilliz
config:
@@ -129,7 +153,33 @@ _Coming soon_
## Pinecone
_Coming soon_
Install pinecone related dependencies using the following command:
```bash
pip install --upgrade 'embedchain[pinecone]'
```
In order to use Pinecone as vector database, set the environment variables `PINECONE_API_KEY` and `PINECONE_ENV` which you can find on [Pinecone dashboard](https://app.pinecone.io/).
<CodeGroup>
```python main.py
from embedchain import App
# load pinecone configuration from yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml config.yaml
vectordb:
provider: pinecone
config:
metric: cosine
vector_dimension: 1536
collection_name: my-pinecone-index
```
</CodeGroup>
## Qdrant
@@ -137,6 +187,21 @@ _Coming soon_
## Weaviate
_Coming soon_
In order to use Weaviate as a vector database, set the environment variables `WEAVIATE_ENDPOINT` and `WEAVIATE_API_KEY` which you can find on [Weaviate dashboard](https://console.weaviate.cloud/dashboard).
```python main.py
from embedchain import App
# load weaviate configuration from yaml file
app = App.from_config(yaml_path="config.yaml")
```
```yaml config.yaml
vectordb:
provider: weaviate
config:
collection_name: my_weaviate_index
```
<Snippet file="missing-vector-db-tip.mdx" />
+3 -6
View File
@@ -35,12 +35,9 @@ embedchain is built on the following stack:
## Team
### Author
### Authors
- Taranjeet Singh ([@taranjeetio](https://twitter.com/taranjeetio))
### Maintainer
- Deshraj Yadav ([@deshrajdry](https://twitter.com/taranjeetio))
### Citation
@@ -49,8 +46,8 @@ If you utilize this repository, please consider citing it with:
```
@misc{embedchain,
author = {Taranjeet Singh},
title = {Embechain: Framework to easily create LLM powered bots over any dataset},
author = {Taranjeet Singh, Deshraj Yadav},
title = {Embechain: Data platform for LLMs - Load, index, retrieve and sync any unstructured data},
year = {2023},
publisher = {GitHub},
journal = {GitHub repository},
+1 -4
View File
@@ -1,6 +1,6 @@
---
title: 📚 Introduction
description: '📝 Embedchain is a framework to easily create LLM powered apps on your data.'
description: '📝 Embedchain is a Data Platform for LLMs - load, index, retrieve, and sync any unstructured data'
---
## 🤔 What is Embedchain?
@@ -27,9 +27,6 @@ naval_bot.add(("Who is Naval Ravikant?", "Naval Ravikant is an Indian-American e
naval_bot.query("What unique capacity does Naval argue humans possess when it comes to understanding explanations or concepts?")
# Answer: Naval argues that humans possess the unique capacity to understand explanations or concepts to the maximum extent possible in this physical reality.
# Ask questions with specific context
naval_bot.query("What unique capacity does Naval argue humans possess when it comes to understanding explanations or concepts?", where={'chapter': 'philosophy'})
```
## 🚀 How it works?
+8 -2
View File
@@ -3,6 +3,8 @@ title: '🚀 Quickstart'
description: '💡 Start building LLM powered apps under 30 seconds'
---
Embedchain is a Data Platform for LLMs - load, index, retrieve, and sync any unstructured data. Using embedchain, you can easily create LLM powered apps over any data.
Install embedchain python package:
```bash
@@ -20,9 +22,11 @@ app = App()
</Step>
<Step title="🗃️ Add data sources">
```python
# Embed online resources
# Add different data sources
elon_bot.add("https://en.wikipedia.org/wiki/Elon_Musk")
elon_bot.add("https://www.forbes.com/profile/elon-musk")
# You can also add local data sources such as pdf, csv files etc.
# elon_bot.add("/path/to/file.pdf")
```
</Step>
<Step title="💬 Query or chat on your data and get answers">
@@ -42,9 +46,11 @@ from embedchain import App
os.environ["OPENAI_API_KEY"] = "xxx"
elon_bot = App()
# Embed online resources
# Add different data sources
elon_bot.add("https://en.wikipedia.org/wiki/Elon_Musk")
elon_bot.add("https://www.forbes.com/profile/elon-musk")
# You can also add local data sources such as pdf, csv files etc.
# elon_bot.add("/path/to/file.pdf")
response = elon_bot.query("What is the net worth of Elon Musk today?")
print(response)
+5 -6
View File
@@ -2,8 +2,7 @@ from string import Template
from embedchain.apps.app import App
from embedchain.apps.open_source_app import OpenSourceApp
from embedchain.config import BaseLlmConfig
from embedchain.config.apps.base_app_config import BaseAppConfig
from embedchain.config import AppConfig, BaseLlmConfig
from embedchain.config.llm.base import (DEFAULT_PROMPT,
DEFAULT_PROMPT_WITH_HISTORY)
from embedchain.helper.json_serializable import register_deserializable
@@ -16,16 +15,16 @@ class EmbedChainPersonApp:
This bot behaves and speaks like a person.
:param person: name of the person, better if its a well known person.
:param config: BaseAppConfig instance to load as configuration.
:param config: AppConfig instance to load as configuration.
"""
def __init__(self, person: str, config: BaseAppConfig = None):
def __init__(self, person: str, config: AppConfig = None):
"""Initialize a new person app
:param person: Name of the person that's imitated.
:type person: str
:param config: Configuration class instance, defaults to None
:type config: BaseAppConfig, optional
:type config: AppConfig, optional
"""
self.person = person
self.person_prompt = f"You are {person}. Whatever you say, you will always say in {person} style." # noqa:E501
@@ -70,7 +69,7 @@ class PersonApp(EmbedChainPersonApp, App):
"""
def query(self, input_query, config: BaseLlmConfig = None, dry_run=False):
config = self.add_person_template_to_config(DEFAULT_PROMPT, config, where=None)
config = self.add_person_template_to_config(DEFAULT_PROMPT, config)
return super().query(input_query, config, dry_run, where=None)
def chat(self, input_query, config: BaseLlmConfig = None, dry_run=False, where=None):
+20
View File
@@ -0,0 +1,20 @@
from typing import Dict, Optional
from embedchain.config.vectordb.base import BaseVectorDbConfig
from embedchain.helper.json_serializable import register_deserializable
@register_deserializable
class PineconeDBConfig(BaseVectorDbConfig):
def __init__(
self,
collection_name: Optional[str] = None,
dir: Optional[str] = None,
vector_dimension: int = 1536,
metric: Optional[str] = "cosine",
**extra_params: Dict[str, any],
):
self.metric = metric
self.vector_dimension = vector_dimension
self.extra_params = extra_params
super().__init__(collection_name=collection_name, dir=dir)
+16
View File
@@ -0,0 +1,16 @@
from typing import Dict, Optional
from embedchain.config.vectordb.base import BaseVectorDbConfig
from embedchain.helper.json_serializable import register_deserializable
@register_deserializable
class WeaviateDBConfig(BaseVectorDbConfig):
def __init__(
self,
collection_name: Optional[str] = None,
dir: Optional[str] = None,
**extra_params: Dict[str, any],
):
self.extra_params = extra_params
super().__init__(collection_name=collection_name, dir=dir)
+5 -4
View File
@@ -339,7 +339,6 @@ class EmbedChain(JSONSerializable):
metadatas = embeddings_data["metadatas"]
ids = embeddings_data["ids"]
new_doc_id = embeddings_data["doc_id"]
if existing_doc_id and existing_doc_id == new_doc_id:
print("Doc content has not changed. Skipping creating chunks and embeddings")
return [], [], [], 0
@@ -351,12 +350,15 @@ class EmbedChain(JSONSerializable):
# get existing ids, and discard doc if any common id exist.
where = {"url": src}
# if data type is qna_pair, we check for question
if chunker.data_type == DataType.QNA_PAIR:
where = {"question": src[0]}
if self.config.id is not None:
where["app_id"] = self.config.id
db_result = self.db.get(ids=ids, where=where) # optional filter
existing_ids = set(db_result["ids"])
if len(existing_ids):
data_dict = {id: (doc, meta) for id, doc, meta in zip(ids, documents, metadatas)}
data_dict = {id: value for id, value in data_dict.items() if id not in existing_ids}
@@ -404,6 +406,7 @@ class EmbedChain(JSONSerializable):
skip_embedding=(chunker.data_type == DataType.IMAGES),
)
count_new_chunks = self.db.count() - chunks_before_addition
print((f"Successfully saved {src} ({chunker.data_type}). New chunks count: {count_new_chunks}"))
return list(documents), metadatas, ids, count_new_chunks
@@ -432,7 +435,6 @@ class EmbedChain(JSONSerializable):
:rtype: List[str]
"""
query_config = config or self.llm.config
if where is not None:
where = where
elif query_config is not None and query_config.where is not None:
@@ -459,7 +461,6 @@ class EmbedChain(JSONSerializable):
where=where,
skip_embedding=(hasattr(config, "query_type") and config.query_type == "Images"),
)
return contents
def query(self, input_query: str, config: BaseLlmConfig = None, dry_run=False, where: Optional[Dict] = None) -> str:
+6
View File
@@ -44,10 +44,12 @@ class EmbedderFactory:
"gpt4all": "embedchain.embedder.gpt4all.GPT4AllEmbedder",
"huggingface": "embedchain.embedder.huggingface.HuggingFaceEmbedder",
"vertexai": "embedchain.embedder.vertexai.VertexAIEmbedder",
"azure_openai": "embedchain.embedder.openai.OpenAIEmbedder",
"openai": "embedchain.embedder.openai.OpenAIEmbedder",
}
provider_to_config_class = {
"openai": "embedchain.config.embedder.base.BaseEmbedderConfig",
"azure_openai": "embedchain.config.embedder.base.BaseEmbedderConfig",
}
@classmethod
@@ -69,11 +71,15 @@ class VectorDBFactory:
"chroma": "embedchain.vectordb.chroma.ChromaDB",
"elasticsearch": "embedchain.vectordb.elasticsearch.ElasticsearchDB",
"opensearch": "embedchain.vectordb.opensearch.OpenSearchDB",
"pinecone": "embedchain.vectordb.pinecone.PineconeDB",
"weaviate": "embedchain.vectordb.weaviate.WeaviateDB",
}
provider_to_config_class = {
"chroma": "embedchain.config.vectordb.chroma.ChromaDbConfig",
"elasticsearch": "embedchain.config.vectordb.elasticsearch.ElasticsearchDBConfig",
"opensearch": "embedchain.config.vectordb.opensearch.OpenSearchDBConfig",
"pinecone": "embedchain.config.vectordb.pinecone.PineconeDBConfig",
"weaviate": "embedchain.config.vectordb.weaviate.WeaviateDBConfig",
}
@classmethod
+6 -3
View File
@@ -129,8 +129,12 @@ class BaseLlm(JSONSerializable):
:return: Search results
:rtype: Unknown
"""
from langchain.tools import DuckDuckGoSearchRun
try:
from langchain.tools import DuckDuckGoSearchRun
except ImportError:
raise ImportError(
'Searching requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
search = DuckDuckGoSearchRun()
logging.info(f"Access search to get answers for {input_query}")
return search.run(input_query)
@@ -202,7 +206,6 @@ class BaseLlm(JSONSerializable):
k["web_search_result"] = self.access_search_and_get_results(input_query)
prompt = self.generate_prompt(input_query, contexts, **k)
logging.info(f"Prompt: {prompt}")
if dry_run:
return prompt
+4 -9
View File
@@ -13,13 +13,9 @@ class OpenAILlm(BaseLlm):
def __init__(self, config: Optional[BaseLlmConfig] = None):
super().__init__(config=config)
def get_llm_model_answer(self, prompt):
def get_llm_model_answer(self, prompt) -> str:
response = OpenAILlm._get_answer(prompt, self.config)
if self.config.stream:
return response
else:
return response.content
return response
def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
messages = []
@@ -35,10 +31,9 @@ class OpenAILlm(BaseLlm):
if config.top_p:
kwargs["model_kwargs"]["top_p"] = config.top_p
if config.stream:
from langchain.callbacks.streaming_stdout import \
StreamingStdOutCallbackHandler
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()])
else:
chat = ChatOpenAI(**kwargs)
return chat(messages)
return chat(messages).content
+8 -1
View File
@@ -3,7 +3,14 @@ import logging
from urllib.parse import urljoin, urlparse
import requests
from bs4 import BeautifulSoup
try:
from bs4 import BeautifulSoup
except ImportError:
raise ImportError(
'DocsSite requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
+6 -2
View File
@@ -1,7 +1,11 @@
import hashlib
from langchain.document_loaders import Docx2txtLoader
try:
from langchain.document_loaders import Docx2txtLoader
except ImportError:
raise ImportError(
'Docx file requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
+6 -2
View File
@@ -1,7 +1,11 @@
import hashlib
from langchain.document_loaders import PyPDFLoader
try:
from langchain.document_loaders import PyPDFLoader
except ImportError:
raise ImportError(
'PDF File requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
from embedchain.utils import clean_string
+8 -2
View File
@@ -2,8 +2,14 @@ import hashlib
import logging
import requests
from bs4 import BeautifulSoup
from bs4.builder import ParserRejectedMarkup
try:
from bs4 import BeautifulSoup
from bs4.builder import ParserRejectedMarkup
except ImportError:
raise ImportError(
'Sitemap requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
+7 -1
View File
@@ -2,7 +2,13 @@ import hashlib
import logging
import requests
from bs4 import BeautifulSoup
try:
from bs4 import BeautifulSoup
except ImportError:
raise ImportError(
'Webpage requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
+6 -2
View File
@@ -1,7 +1,11 @@
import hashlib
from langchain.document_loaders import UnstructuredXMLLoader
try:
from langchain.document_loaders import UnstructuredXMLLoader
except ImportError:
raise ImportError(
'XML file requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
from embedchain.utils import clean_string
+6 -2
View File
@@ -1,7 +1,11 @@
import hashlib
from langchain.document_loaders import YoutubeLoader
try:
from langchain.document_loaders import YoutubeLoader
except ImportError:
raise ImportError(
'YouTube video requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`'
) from None
from embedchain.helper.json_serializable import register_deserializable
from embedchain.loaders.base_loader import BaseLoader
from embedchain.utils import clean_string
+2 -4
View File
@@ -1,5 +1,5 @@
import logging
from typing import Any, Dict, List, Optional, Set
from typing import Any, Dict, List, Optional
try:
from elasticsearch import Elasticsearch
@@ -74,9 +74,7 @@ class ElasticsearchDB(BaseVectorDB):
def _get_or_create_collection(self, name):
"""Note: nothing to return here. Discuss later"""
def get(
self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None
) -> Set[str]:
def get(self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None):
"""
Get existing doc ids present in vector database
+182
View File
@@ -0,0 +1,182 @@
import os
from typing import Dict, List, Optional
try:
import pinecone
except ImportError:
raise ImportError(
"Pinecone requires extra dependencies. Install with `pip install --upgrade 'embedchain[pinecone]'`"
) from None
from embedchain.config.vectordb.pinecone import PineconeDBConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.vectordb.base import BaseVectorDB
@register_deserializable
class PineconeDB(BaseVectorDB):
"""
Pinecone as vector database
"""
BATCH_SIZE = 100
def __init__(
self,
config: Optional[PineconeDBConfig] = None,
):
"""Pinecone as vector database.
:param config: Pinecone database config, defaults to None
:type config: PineconeDBConfig, optional
:raises ValueError: No config provided
"""
if config is None:
self.config = PineconeDBConfig()
else:
if not isinstance(config, PineconeDBConfig):
raise TypeError(
"config is not a `PineconeDBConfig` instance. "
"Please make sure the type is right and that you are passing an instance."
)
self.config = config
self.client = self._setup_pinecone_index()
# Call parent init here because embedder is needed
super().__init__(config=self.config)
def _initialize(self):
"""
This method is needed because `embedder` attribute needs to be set externally before it can be initialized.
"""
if not self.embedder:
raise ValueError("Embedder not set. Please set an embedder with `set_embedder` before initialization.")
# Loads the Pinecone index or creates it if not present.
def _setup_pinecone_index(self):
pinecone.init(
api_key=os.environ.get("PINECONE_API_KEY"),
environment=os.environ.get("PINECONE_ENV"),
**self.config.extra_params,
)
self.index_name = self._get_index_name()
indexes = pinecone.list_indexes()
if indexes is None or self.index_name not in indexes:
pinecone.create_index(
name=self.index_name, metric=self.config.metric, dimension=self.config.vector_dimension
)
return pinecone.Index(self.index_name)
def get(self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None):
"""
Get existing doc ids present in vector database
:param ids: _list of doc ids to check for existence
:type ids: List[str]
:param where: to filter data
:type where: Dict[str, any]
:return: ids
:rtype: Set[str]
"""
existing_ids = list()
if ids is not None:
for i in range(0, len(ids), 1000):
result = self.client.fetch(ids=ids[i : i + 1000])
batch_existing_ids = list(result.get("vectors").keys())
existing_ids.extend(batch_existing_ids)
return {"ids": existing_ids}
def add(
self,
embeddings: List[List[float]],
documents: List[str],
metadatas: List[object],
ids: List[str],
skip_embedding: bool,
):
"""add data in vector database
:param documents: list of texts to add
:type documents: List[str]
:param metadatas: list of metadata associated with docs
:type metadatas: List[object]
:param ids: ids of docs
:type ids: List[str]
"""
docs = []
print("Adding documents to Pinecone...")
embeddings = self.embedder.embedding_fn(documents)
for id, text, metadata, embedding in zip(ids, documents, metadatas, embeddings):
docs.append(
{
"id": id,
"values": embedding,
"metadata": {**metadata, "text": text},
}
)
for i in range(0, len(docs), self.BATCH_SIZE):
self.client.upsert(docs[i : i + self.BATCH_SIZE])
def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
"""
query contents from vector database based on vector similarity
:param input_query: list of query string
:type input_query: List[str]
:param n_results: no of similar documents to fetch from database
:type n_results: int
:param where: Optional. to filter data
:type where: Dict[str, any]
:param skip_embedding: Optional. if True, input_query is already embedded
:type skip_embedding: bool
:return: Database contents that are the result of the query
:rtype: List[str]
"""
if not skip_embedding:
query_vector = self.embedder.embedding_fn([input_query])[0]
else:
query_vector = input_query
contents = self.client.query(vector=query_vector, filter=where, top_k=n_results, include_metadata=True)
embeddings = list(map(lambda content: content["metadata"]["text"], contents["matches"]))
return embeddings
def set_collection_name(self, name: str):
"""
Set the name of the collection. A collection is an isolated space for vectors.
:param name: Name of the collection.
:type name: str
"""
if not isinstance(name, str):
raise TypeError("Collection name must be a string")
self.config.collection_name = name
def count(self) -> int:
"""
Count number of documents/chunks embedded in the database.
:return: number of documents
:rtype: int
"""
return self.client.describe_index_stats()["total_vector_count"]
def _get_or_create_db(self):
"""Called during initialization"""
return self.client
def reset(self):
"""
Resets the database. Deletes all embeddings irreversibly.
"""
# Delete all data from the database
pinecone.delete_index(self.index_name)
self._setup_pinecone_index()
# Pinecone only allows alphanumeric characters and "-" in the index name
def _get_index_name(self) -> str:
"""Get the Pinecone index for a collection
:return: Pinecone index
:rtype: str
"""
return f"{self.config.collection_name}-{self.config.vector_dimension}".lower().replace("_", "-")
+297
View File
@@ -0,0 +1,297 @@
import copy
import os
from typing import Dict, List, Optional
try:
import weaviate
except ImportError:
raise ImportError(
"Weaviate requires extra dependencies. Install with `pip install --upgrade 'embedchain[weaviate]'`"
) from None
from embedchain.config.vectordb.weaviate import WeaviateDBConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.vectordb.base import BaseVectorDB
@register_deserializable
class WeaviateDB(BaseVectorDB):
"""
Weaviate as vector database
"""
BATCH_SIZE = 100
def __init__(
self,
config: Optional[WeaviateDBConfig] = None,
):
"""Weaviate as vector database.
:param config: Weaviate database config, defaults to None
:type config: WeaviateDBConfig, optional
:raises ValueError: No config provided
"""
if config is None:
self.config = WeaviateDBConfig()
else:
if not isinstance(config, WeaviateDBConfig):
raise TypeError(
"config is not a `WeaviateDBConfig` instance. "
"Please make sure the type is right and that you are passing an instance."
)
self.config = config
self.client = weaviate.Client(
url=os.environ.get("WEAVIATE_ENDPOINT"),
auth_client_secret=weaviate.AuthApiKey(api_key=os.environ.get("WEAVIATE_API_KEY")),
**self.config.extra_params,
)
# Call parent init here because embedder is needed
super().__init__(config=self.config)
def _initialize(self):
"""
This method is needed because `embedder` attribute needs to be set externally before it can be initialized.
"""
if not self.embedder:
raise ValueError("Embedder not set. Please set an embedder with `set_embedder` before initialization.")
self.index_name = self._get_index_name()
self.metadata_keys = {"data_type", "doc_id", "url", "hash", "app_id", "text"}
if not self.client.schema.exists(self.index_name):
# id is a reserved field in Weaviate, hence we had to change the name of the id field to identifier
# The none vectorizer is crucial as we have our own custom embedding function
class_obj = {
"classes": [
{
"class": self.index_name,
"vectorizer": "none",
"properties": [
{
"name": "identifier",
"dataType": ["text"],
},
{
"name": "text",
"dataType": ["text"],
},
{
"name": "metadata",
"dataType": [self.index_name + "_metadata"],
},
],
},
{
"class": self.index_name + "_metadata",
"vectorizer": "none",
"properties": [
{
"name": "data_type",
"dataType": ["text"],
},
{
"name": "doc_id",
"dataType": ["text"],
},
{
"name": "url",
"dataType": ["text"],
},
{
"name": "hash",
"dataType": ["text"],
},
{
"name": "app_id",
"dataType": ["text"],
},
{
"name": "text",
"dataType": ["text"],
},
],
},
]
}
self.client.schema.create(class_obj)
def get(self, ids: Optional[List[str]] = None, where: Optional[Dict[str, any]] = None, limit: Optional[int] = None):
"""
Get existing doc ids present in vector database
:param ids: _list of doc ids to check for existance
:type ids: List[str]
:param where: to filter data
:type where: Dict[str, any]
:return: ids
:rtype: Set[str]
"""
if ids is None or len(ids) == 0:
return {"ids": []}
existing_ids = []
cursor = None
has_iterated_once = False
while cursor is not None or not has_iterated_once:
has_iterated_once = True
results = self._query_with_cursor(
self.client.query.get(self.index_name, ["identifier"])
.with_additional(["id"])
.with_limit(self.BATCH_SIZE),
cursor,
)
fetched_results = results["data"]["Get"].get(self.index_name, [])
if len(fetched_results) == 0:
break
for result in fetched_results:
existing_ids.append(result["identifier"])
cursor = result["_additional"]["id"]
return {"ids": existing_ids}
def add(
self,
embeddings: List[List[float]],
documents: List[str],
metadatas: List[object],
ids: List[str],
skip_embedding: bool,
):
"""add data in vector database
:param embeddings: list of embeddings for the corresponding documents to be added
:type documents: List[List[float]]
:param documents: list of texts to add
:type documents: List[str]
:param metadatas: list of metadata associated with docs
:type metadatas: List[object]
:param ids: ids of docs
:type ids: List[str]
:param skip_embedding: A boolean flag indicating if the embedding for the documents to be added is to be
generated or not
:type skip_embedding: bool
"""
print("Adding documents to Weaviate...")
if not skip_embedding:
embeddings = self.embedder.embedding_fn(documents)
self.client.batch.configure(batch_size=self.BATCH_SIZE, timeout_retries=3) # Configure batch
with self.client.batch as batch: # Initialize a batch process
for id, text, metadata, embedding in zip(ids, documents, metadatas, embeddings):
doc = {"identifier": id, "text": text}
updated_metadata = {"text": text}
if metadata is not None:
updated_metadata.update(**metadata)
obj_uuid = batch.add_data_object(
data_object=copy.deepcopy(doc), class_name=self.index_name, vector=embedding
)
metadata_uuid = batch.add_data_object(
data_object=copy.deepcopy(updated_metadata),
class_name=self.index_name + "_metadata",
vector=embedding,
)
batch.add_reference(obj_uuid, self.index_name, "metadata", metadata_uuid, self.index_name + "_metadata")
def query(self, input_query: List[str], n_results: int, where: Dict[str, any], skip_embedding: bool) -> List[str]:
"""
query contents from vector database based on vector similarity
:param input_query: list of query string
:type input_query: List[str]
:param n_results: no of similar documents to fetch from database
:type n_results: int
:param where: Optional. to filter data
:type where: Dict[str, any]
:param skip_embedding: A boolean flag indicating if the embedding for the documents to be added is to be
generated or not
:type skip_embedding: bool
:return: Database contents that are the result of the query
:rtype: List[str]
"""
if not skip_embedding:
query_vector = self.embedder.embedding_fn([input_query])[0]
else:
query_vector = input_query
keys = set(where.keys() if where is not None else set())
if len(keys.intersection(self.metadata_keys)) != 0:
weaviate_where_operands = []
for key in keys:
if key in self.metadata_keys:
weaviate_where_operands.append(
{
"path": ["metadata", self.index_name + "_metadata", key],
"operator": "Equal",
"valueText": where.get(key),
}
)
if len(weaviate_where_operands) == 1:
weaviate_where_clause = weaviate_where_operands[0]
else:
weaviate_where_clause = {"operator": "And", "operands": weaviate_where_operands}
results = (
self.client.query.get(self.index_name, ["text"])
.with_where(weaviate_where_clause)
.with_near_vector({"vector": query_vector})
.with_limit(n_results)
.do()
)
else:
results = (
self.client.query.get(self.index_name, ["text"])
.with_near_vector({"vector": query_vector})
.with_limit(n_results)
.do()
)
matched_tokens = []
for result in results["data"]["Get"].get(self.index_name):
matched_tokens.append(result["text"])
return matched_tokens
def set_collection_name(self, name: str):
"""
Set the name of the collection. A collection is an isolated space for vectors.
:param name: Name of the collection.
:type name: str
"""
if not isinstance(name, str):
raise TypeError("Collection name must be a string")
self.config.collection_name = name
def count(self) -> int:
"""
Count number of documents/chunks embedded in the database.
:return: number of documents
:rtype: int
"""
data = self.client.query.aggregate(self.index_name).with_meta_count().do()
return data["data"]["Aggregate"].get(self.index_name)[0]["meta"]["count"]
def _get_or_create_db(self):
"""Called during initialization"""
return self.client
def reset(self):
"""
Resets the database. Deletes all embeddings irreversibly.
"""
# Delete all data from the database
self.client.batch.delete_objects(
self.index_name, where={"path": ["identifier"], "operator": "Like", "valueText": ".*"}
)
# Weaviate internally by default capitalizes the class name
def _get_index_name(self) -> str:
"""Get the Weaviate index for a collection
:return: Weaviate index
:rtype: str
"""
return f"{self.config.collection_name}_{self.embedder.vector_dimension}".capitalize()
def _query_with_cursor(self, query, cursor):
if cursor is not None:
query.with_after(cursor)
results = query.do()
return results
@@ -28,7 +28,7 @@ export default function Home() {
Welcome to Embedchain Playground
</h1>
<p className="mb-6 text-lg font-normal text-gray-500 lg:text-xl">
embedchain is a framework to easily create LLM powered bots over any
Embedchain is a Data Platform for LLMs - Load, index, retrieve, and sync any unstructured data
dataset
</p>
</div>
+182
View File
@@ -0,0 +1,182 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "63ab5e89",
"metadata": {},
"source": [
"## Cookbook for using Azure OpenAI with Embedchain"
]
},
{
"cell_type": "markdown",
"id": "e32a0265",
"metadata": {},
"source": [
"### Step-1: Install embedchain package"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b80ff15a",
"metadata": {},
"outputs": [],
"source": [
"!pip install embedchain"
]
},
{
"cell_type": "markdown",
"id": "ac982a56",
"metadata": {},
"source": [
"### Step-2: Set Azure OpenAI related environment variables\n",
"\n",
"You can find these env variables on your Azure OpenAI dashboard."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e0a36133",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from embedchain import App\n",
"\n",
"os.environ[\"OPENAI_API_TYPE\"] = \"azure\"\n",
"os.environ[\"OPENAI_API_BASE\"] = \"https://xxx.openai.azure.com/\"\n",
"os.environ[\"OPENAI_API_KEY\"] = \"xxx\"\n",
"os.environ[\"OPENAI_API_VERSION\"] = \"xxx\""
]
},
{
"cell_type": "markdown",
"id": "7d7b554e",
"metadata": {},
"source": [
"### Step-3: Define your llm and embedding model config"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b9f52fc5",
"metadata": {},
"outputs": [],
"source": [
"config = \"\"\"\n",
"llm:\n",
" provider: azure_openai\n",
" model: gpt-35-turbo\n",
" config:\n",
" deployment_name: ec_openai_azure\n",
" temperature: 0.5\n",
" max_tokens: 1000\n",
" top_p: 1\n",
" stream: false\n",
"\n",
"embedder:\n",
" provider: azure_openai\n",
" config:\n",
" model: text-embedding-ada-002\n",
" deployment_name: ec_embeddings_ada_002\n",
"\"\"\"\n",
"\n",
"# Write the multi-line string to a YAML file\n",
"with open('azure_openai.yaml', 'w') as file:\n",
" file.write(config)"
]
},
{
"cell_type": "markdown",
"id": "98a11130",
"metadata": {},
"source": [
"### Step-4 Create embedchain app based on the config"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1ee9bdd9",
"metadata": {},
"outputs": [],
"source": [
"app = App.from_config(yaml_path=\"azure_openai.yaml\")"
]
},
{
"cell_type": "markdown",
"id": "554dc97b",
"metadata": {},
"source": [
"### Step-5: Add data sources to your app"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "686ae765",
"metadata": {},
"outputs": [],
"source": [
"app.add(\"https://www.forbes.com/profile/elon-musk\")"
]
},
{
"cell_type": "markdown",
"id": "ccc7d421",
"metadata": {},
"source": [
"### Step-6: All set. Now start asking questions related to your data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "27868a7d",
"metadata": {},
"outputs": [],
"source": [
"while(True):\n",
" question = input(\"Enter question: \")\n",
" if question in ['q', 'exit', 'quit']\n",
" break\n",
" answer = app.query(question)\n",
" print(answer)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1f2ead5",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.4"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
+27 -11
View File
@@ -1,7 +1,7 @@
[tool.poetry]
name = "embedchain"
version = "0.0.70"
description = "Embedchain is a framework to easily create LLM powered apps over any dataset"
version = "0.0.73"
description = "Data platform for LLMs - Load, index, retrieve and sync any unstructured data"
authors = ["Taranjeet Singh, Deshraj Yadav"]
license = "Apache License"
readme = "README.md"
@@ -90,13 +90,13 @@ python-dotenv = "^1.0.0"
langchain = "^0.0.279"
requests = "^2.31.0"
openai = "^0.27.5"
tiktoken = "^0.4.0"
tiktoken = { version="^0.4.0", optional=true }
chromadb ="^0.4.8"
youtube-transcript-api = "^0.6.1"
beautifulsoup4 = "^4.12.2"
pypdf = "^3.11.0"
pytube = "^15.0.0"
duckduckgo-search = "^3.8.5"
youtube-transcript-api = { version="^0.6.1", optional=true }
beautifulsoup4 = { version="^4.12.2", optional=true }
pypdf = { version="^3.11.0", optional=true }
pytube = { version="^15.0.0", optional=true }
duckduckgo-search = { version="^3.8.5", optional=true }
llama-hub = { version = "^0.0.29", optional = true }
sentence-transformers = { version = "^2.2.2", optional = true }
torch = { version = "2.0.0", optional = true }
@@ -111,14 +111,16 @@ fastapi-poe = { version = "0.0.16", optional = true }
discord = { version = "^2.3.2", optional = true }
slack-sdk = { version = "3.21.3", optional = true }
cohere = { version = "^4.27", optional= true }
docx2txt = "^0.8"
unstructured = {extras = ["local-inference"], version = "^0.10.18"}
weaviate-client = { version = "^3.24.1", optional= true }
docx2txt = { version="^0.8", optional=true }
pinecone-client = { version = "^2.2.4", optional = true }
unstructured = {extras = ["local-inference"], version = "^0.10.18", optional=true}
pillow = { version = "10.0.1", optional = true }
torchvision = { version = ">=0.15.1, !=0.15.2", optional = true }
ftfy = { version = "6.1.1", optional = true }
regex = { version = "2023.8.8", optional = true }
huggingface_hub = { version = "^0.17.3", optional = true }
pymilvus = "2.3.1"
pymilvus = { version="2.3.1", optional = true }
[tool.poetry.group.dev.dependencies]
black = "^23.3.0"
@@ -131,6 +133,8 @@ click = "^8.1.3"
isort = "^5.12.0"
pytest-cov = "^4.1.0"
responses = "^0.23.3"
mock = "^5.1.0"
pytest-asyncio = "^0.21.1"
[tool.poetry.extras]
streamlit = ["streamlit"]
@@ -142,10 +146,22 @@ poe = ["fastapi-poe"]
discord = ["discord"]
slack = ["slack-sdk", "flask"]
whatsapp = ["twilio", "flask"]
weaviate = ["weaviate-client"]
pinecone = ["pinecone-client"]
images = ["torch", "ftfy", "regex", "pillow", "torchvision"]
huggingface_hub=["huggingface_hub"]
cohere = ["cohere"]
milvus = ["pymilvus"]
dataloaders=[
"youtube-transcripts-api",
"beautifulsoup4",
"pypdf",
"pytube",
"duckduckgo-search",
"docx2txt",
"unstructured",
"sentence-transformers",
]
[tool.poetry.group.docs.dependencies]
+77 -64
View File
@@ -1,6 +1,6 @@
import os
import unittest
import pytest
import yaml
from embedchain import App, CustomApp, Llama2App, OpenSourceApp
@@ -12,97 +12,110 @@ from embedchain.vectordb.base import BaseVectorDB, BaseVectorDbConfig
from embedchain.vectordb.chroma import ChromaDB
class TestApps(unittest.TestCase):
@pytest.fixture
def app():
os.environ["OPENAI_API_KEY"] = "test_api_key"
def test_app(self):
app = App()
self.assertIsInstance(app.llm, BaseLlm)
self.assertIsInstance(app.db, BaseVectorDB)
self.assertIsInstance(app.embedder, BaseEmbedder)
wrong_llm = "wrong_llm"
with self.assertRaises(TypeError):
App(llm=wrong_llm)
wrong_db = "wrong_db"
with self.assertRaises(TypeError):
App(db=wrong_db)
wrong_embedder = "wrong_embedder"
with self.assertRaises(TypeError):
App(embedder=wrong_embedder)
def test_custom_app(self):
app = CustomApp()
self.assertIsInstance(app.llm, BaseLlm)
self.assertIsInstance(app.db, BaseVectorDB)
self.assertIsInstance(app.embedder, BaseEmbedder)
def test_opensource_app(self):
app = OpenSourceApp()
self.assertIsInstance(app.llm, BaseLlm)
self.assertIsInstance(app.db, BaseVectorDB)
self.assertIsInstance(app.embedder, BaseEmbedder)
def test_llama2_app(self):
os.environ["REPLICATE_API_TOKEN"] = "-"
app = Llama2App()
self.assertIsInstance(app.llm, BaseLlm)
self.assertIsInstance(app.db, BaseVectorDB)
self.assertIsInstance(app.embedder, BaseEmbedder)
return App()
class TestConfigForAppComponents(unittest.TestCase):
collection_name = "my-test-collection"
@pytest.fixture
def custom_app():
os.environ["OPENAI_API_KEY"] = "test_api_key"
return CustomApp()
@pytest.fixture
def opensource_app():
os.environ["OPENAI_API_KEY"] = "test_api_key"
return OpenSourceApp()
@pytest.fixture
def llama2_app():
os.environ["OPENAI_API_KEY"] = "test_api_key"
os.environ["REPLICATE_API_TOKEN"] = "-"
return Llama2App()
def test_app(app):
assert isinstance(app.llm, BaseLlm)
assert isinstance(app.db, BaseVectorDB)
assert isinstance(app.embedder, BaseEmbedder)
def test_custom_app(custom_app):
assert isinstance(custom_app.llm, BaseLlm)
assert isinstance(custom_app.db, BaseVectorDB)
assert isinstance(custom_app.embedder, BaseEmbedder)
def test_opensource_app(opensource_app):
assert isinstance(opensource_app.llm, BaseLlm)
assert isinstance(opensource_app.db, BaseVectorDB)
assert isinstance(opensource_app.embedder, BaseEmbedder)
def test_llama2_app(llama2_app):
assert isinstance(llama2_app.llm, BaseLlm)
assert isinstance(llama2_app.db, BaseVectorDB)
assert isinstance(llama2_app.embedder, BaseEmbedder)
class TestConfigForAppComponents:
def test_constructor_config(self):
"""
Test that app can be configured through the app constructor.
"""
app = App(db_config=ChromaDbConfig(collection_name=self.collection_name))
self.assertEqual(app.db.config.collection_name, self.collection_name)
collection_name = "my-test-collection"
app = App(db_config=ChromaDbConfig(collection_name=collection_name))
assert app.db.config.collection_name == collection_name
def test_component_config(self):
"""
Test that app can also be configured by passing a configured component to the app
"""
database = ChromaDB(config=ChromaDbConfig(collection_name=self.collection_name))
collection_name = "my-test-collection"
database = ChromaDB(config=ChromaDbConfig(collection_name=collection_name))
app = App(db=database)
self.assertEqual(app.db.config.collection_name, self.collection_name)
assert app.db.config.collection_name == collection_name
def test_different_configs_are_proper_instances(self):
config = AppConfig()
wrong_app_config = AddConfig()
app_config = AppConfig()
wrong_config = AddConfig()
with pytest.raises(TypeError):
App(config=wrong_config)
with self.assertRaises(TypeError):
App(config=wrong_app_config)
self.assertIsInstance(config, AppConfig)
assert isinstance(app_config, AppConfig)
llm_config = BaseLlmConfig()
wrong_llm_config = "wrong_llm_config"
with self.assertRaises(TypeError):
with pytest.raises(TypeError):
App(llm_config=wrong_llm_config)
self.assertIsInstance(llm_config, BaseLlmConfig)
assert isinstance(llm_config, BaseLlmConfig)
db_config = BaseVectorDbConfig()
wrong_db_config = "wrong_db_config"
with self.assertRaises(TypeError):
with pytest.raises(TypeError):
App(db_config=wrong_db_config)
self.assertIsInstance(db_config, BaseVectorDbConfig)
assert isinstance(db_config, BaseVectorDbConfig)
embedder_config = BaseEmbedderConfig()
wrong_embedder_config = "wrong_embedder_config"
with self.assertRaises(TypeError):
with pytest.raises(TypeError):
App(embedder_config=wrong_embedder_config)
self.assertIsInstance(embedder_config, BaseEmbedderConfig)
assert isinstance(embedder_config, BaseEmbedderConfig)
def test_components_raises_type_error_if_not_proper_instances(self):
wrong_llm = "wrong_llm"
with pytest.raises(TypeError):
App(llm=wrong_llm)
wrong_db = "wrong_db"
with pytest.raises(TypeError):
App(db=wrong_db)
wrong_embedder = "wrong_embedder"
with pytest.raises(TypeError):
App(embedder=wrong_embedder)
class TestAppFromConfig:
+81
View File
@@ -0,0 +1,81 @@
import pytest
from embedchain.apps.app import App
from embedchain.apps.person_app import PersonApp, PersonOpenSourceApp
from embedchain.config import AppConfig, BaseLlmConfig
from embedchain.config.llm.base import DEFAULT_PROMPT
@pytest.fixture
def person_app():
config = AppConfig()
return PersonApp("John Doe", config)
@pytest.fixture
def opensource_person_app():
config = AppConfig()
return PersonOpenSourceApp("John Doe", config)
def test_person_app_initialization(person_app):
assert person_app.person == "John Doe"
assert f"You are {person_app.person}" in person_app.person_prompt
assert isinstance(person_app.config, AppConfig)
def test_person_app_add_person_template_to_config_with_invalid_template():
app = PersonApp("John Doe")
default_prompt = "Input Prompt"
with pytest.raises(ValueError):
# as prompt doesn't contain $context and $query
app.add_person_template_to_config(default_prompt)
def test_person_app_add_person_template_to_config_with_valid_template():
app = PersonApp("John Doe")
config = app.add_person_template_to_config(DEFAULT_PROMPT)
assert (
config.template.template
== f"You are John Doe. Whatever you say, you will always say in John Doe style. {DEFAULT_PROMPT}"
)
def test_person_app_query(mocker, person_app):
input_query = "Hello, how are you?"
config = BaseLlmConfig()
mocker.patch.object(App, "query", return_value="Mocked response")
result = person_app.query(input_query, config)
assert result == "Mocked response"
def test_person_app_chat(mocker, person_app):
input_query = "Hello, how are you?"
config = BaseLlmConfig()
mocker.patch.object(App, "chat", return_value="Mocked chat response")
result = person_app.chat(input_query, config)
assert result == "Mocked chat response"
def test_opensource_person_app_query(mocker, opensource_person_app):
input_query = "Hello, how are you?"
config = BaseLlmConfig()
mocker.patch.object(App, "query", return_value="Mocked response")
result = opensource_person_app.query(input_query, config)
assert result == "Mocked response"
def test_opensource_person_app_chat(mocker, opensource_person_app):
input_query = "Hello, how are you?"
config = BaseLlmConfig()
mocker.patch.object(App, "chat", return_value="Mocked chat response")
result = opensource_person_app.chat(input_query, config)
assert result == "Mocked chat response"
+51
View File
@@ -0,0 +1,51 @@
import argparse
import pytest
from fastapi_poe.types import ProtocolMessage, QueryRequest
from embedchain.bots.poe import PoeBot, start_command
@pytest.fixture
def poe_bot(mocker):
bot = PoeBot()
mocker.patch("fastapi_poe.run")
return bot
@pytest.mark.asyncio
async def test_poe_bot_get_response(poe_bot, mocker):
query = QueryRequest(
version="test",
type="query",
query=[ProtocolMessage(role="system", content="Test content")],
user_id="test_user_id",
conversation_id="test_conversation_id",
message_id="test_message_id",
)
mocker.patch.object(poe_bot.app.llm, "set_history")
response_generator = poe_bot.get_response(query)
await response_generator.__anext__()
poe_bot.app.llm.set_history.assert_called_once()
def test_poe_bot_handle_message(poe_bot, mocker):
mocker.patch.object(poe_bot, "ask_bot", return_value="Answer from the bot")
response_ask = poe_bot.handle_message("What is the answer?")
assert response_ask == "Answer from the bot"
# TODO: This test will fail because the add_data method is commented out.
# mocker.patch.object(poe_bot, 'add_data', return_value="Added data from: some_data")
# response_add = poe_bot.handle_message("/add some_data")
# assert response_add == "Added data from: some_data"
def test_start_command(mocker):
mocker.patch("argparse.ArgumentParser.parse_args", return_value=argparse.Namespace(api_key="test_api_key"))
mocker.patch("embedchain.bots.poe.run")
start_command()
+35 -54
View File
@@ -1,70 +1,51 @@
import os
import unittest
from unittest.mock import MagicMock, patch
import pytest
from embedchain import App
from embedchain.config import AddConfig, AppConfig, ChunkerConfig
from embedchain.models.data_type import DataType
os.environ["OPENAI_API_KEY"] = "test_key"
class TestApp(unittest.TestCase):
os.environ["OPENAI_API_KEY"] = "test_key"
def setUp(self):
self.app = App(config=AppConfig(collect_metrics=False))
@pytest.fixture
def app(mocker):
mocker.patch("chromadb.api.models.Collection.Collection.add")
return App(config=AppConfig(collect_metrics=False))
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_add(self):
"""
This test checks the functionality of the 'add' method in the App class.
It begins by simulating the addition of a web page with a specific URL to the application instance.
The 'add' method is expected to append the input type and URL to the 'user_asks' attribute of the App instance.
By asserting that 'user_asks' is updated correctly after the 'add' method is called, we can confirm that the
method is working as intended.
The Collection.add method from the chromadb library is mocked during this test to isolate the behavior of the
'add' method.
"""
self.app.add("https://example.com", metadata={"meta": "meta-data"})
self.assertEqual(self.app.user_asks, [["https://example.com", "web_page", {"meta": "meta-data"}]])
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_add_sitemap(self):
"""
In addition to the test_add function, this test checks that sitemaps can be added with the correct data type.
"""
self.app.add("https://www.google.com/sitemap.xml", metadata={"meta": "meta-data"})
self.assertEqual(self.app.user_asks, [["https://www.google.com/sitemap.xml", "sitemap", {"meta": "meta-data"}]])
def test_add(app):
app.add("https://example.com", metadata={"meta": "meta-data"})
assert app.user_asks == [["https://example.com", "web_page", {"meta": "meta-data"}]]
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_add_forced_type(self):
"""
Test that you can also force a data_type with `add`.
"""
data_type = "text"
self.app.add("https://example.com", data_type=data_type, metadata={"meta": "meta-data"})
self.assertEqual(self.app.user_asks, [["https://example.com", data_type, {"meta": "meta-data"}]])
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_dry_run(self):
"""
Test that if dry_run == True then data chunks are returned.
"""
def test_add_sitemap(app):
app.add("https://www.google.com/sitemap.xml", metadata={"meta": "meta-data"})
assert app.user_asks == [["https://www.google.com/sitemap.xml", "sitemap", {"meta": "meta-data"}]]
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0)
# We can't test with lorem ipsum because chunks are deduped, so would be recurring characters.
text = """0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"""
result = self.app.add(source=text, config=AddConfig(chunker=chunker_config), dry_run=True)
def test_add_forced_type(app):
data_type = "text"
app.add("https://example.com", data_type=data_type, metadata={"meta": "meta-data"})
assert app.user_asks == [["https://example.com", data_type, {"meta": "meta-data"}]]
chunks = result["chunks"]
metadata = result["metadata"]
count = result["count"]
data_type = result["type"]
self.assertEqual(len(chunks), len(text))
self.assertEqual(count, len(text))
self.assertEqual(data_type, DataType.TEXT)
for item in metadata:
self.assertIsInstance(item, dict)
self.assertIn(item["url"], "local")
self.assertIn(item["data_type"], "text")
def test_dry_run(app):
chunker_config = ChunkerConfig(chunk_size=1, chunk_overlap=0)
text = """0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"""
result = app.add(source=text, config=AddConfig(chunker=chunker_config), dry_run=True)
chunks = result["chunks"]
metadata = result["metadata"]
count = result["count"]
data_type = result["type"]
assert len(chunks) == len(text)
assert count == len(text)
assert data_type == DataType.TEXT
for item in metadata:
assert isinstance(item, dict)
assert "local" in item["url"]
assert "text" in item["data_type"]
+65
View File
@@ -0,0 +1,65 @@
import pytest
from string import Template
from embedchain.llm.base import BaseLlm, BaseLlmConfig
@pytest.fixture
def base_llm():
config = BaseLlmConfig()
return BaseLlm(config=config)
def test_is_get_llm_model_answer_not_implemented(base_llm):
with pytest.raises(NotImplementedError):
base_llm.get_llm_model_answer()
def test_is_stream_bool():
with pytest.raises(ValueError):
config = BaseLlmConfig(stream="test value")
BaseLlm(config=config)
def test_template_string_gets_converted_to_Template_instance():
config = BaseLlmConfig(template="test value $query $context")
llm = BaseLlm(config=config)
assert isinstance(llm.config.template, Template)
def test_is_get_llm_model_answer_implemented():
class TestLlm(BaseLlm):
def get_llm_model_answer(self):
return "Implemented"
config = BaseLlmConfig()
llm = TestLlm(config=config)
assert llm.get_llm_model_answer() == "Implemented"
def test_stream_query_response(base_llm):
answer = ["Chunk1", "Chunk2", "Chunk3"]
result = list(base_llm._stream_query_response(answer))
assert result == answer
def test_stream_chat_response(base_llm):
answer = ["Chunk1", "Chunk2", "Chunk3"]
result = list(base_llm._stream_chat_response(answer))
assert result == answer
def test_append_search_and_context(base_llm):
context = "Context"
web_search_result = "Web Search Result"
result = base_llm._append_search_and_context(context, web_search_result)
expected_result = "Context\nWeb Search Result: Web Search Result"
assert result == expected_result
def test_access_search_and_get_results(base_llm, mocker):
base_llm.access_search_and_get_results = mocker.patch.object(
base_llm, "access_search_and_get_results", return_value="Search Results"
)
input_query = "Test query"
result = base_llm.access_search_and_get_results(input_query)
assert result == "Search Results"
+44 -22
View File
@@ -1,33 +1,55 @@
import os
import unittest
from unittest.mock import patch
import pytest
from embedchain.config import BaseLlmConfig
from embedchain.llm.cohere import CohereLlm
class TestCohereLlm(unittest.TestCase):
def setUp(self):
os.environ["COHERE_API_KEY"] = "test_api_key"
self.config = BaseLlmConfig(model="gptd-instruct-tft", max_tokens=50, temperature=0.7, top_p=0.8)
@pytest.fixture
def cohere_llm_config():
os.environ["COHERE_API_KEY"] = "test_api_key"
config = BaseLlmConfig(model="gptd-instruct-tft", max_tokens=50, temperature=0.7, top_p=0.8)
yield config
os.environ.pop("COHERE_API_KEY")
def test_init_raises_value_error_without_api_key(self):
os.environ.pop("COHERE_API_KEY")
with self.assertRaises(ValueError):
CohereLlm()
def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
llm = CohereLlm(self.config)
llm.config.system_prompt = "system_prompt"
with self.assertRaises(ValueError):
llm.get_llm_model_answer("prompt")
def test_init_raises_value_error_without_api_key(mocker):
mocker.patch.dict(os.environ, clear=True)
with pytest.raises(ValueError):
CohereLlm()
@patch("embedchain.llm.cohere.CohereLlm._get_answer")
def test_get_llm_model_answer(self, mock_get_answer):
mock_get_answer.return_value = "Test answer"
llm = CohereLlm(self.config)
answer = llm.get_llm_model_answer("Test query")
def test_get_llm_model_answer_raises_value_error_for_system_prompt(cohere_llm_config):
llm = CohereLlm(cohere_llm_config)
llm.config.system_prompt = "system_prompt"
with pytest.raises(ValueError):
llm.get_llm_model_answer("prompt")
self.assertEqual(answer, "Test answer")
mock_get_answer.assert_called_once()
def test_get_llm_model_answer(cohere_llm_config, mocker):
mocker.patch("embedchain.llm.cohere.CohereLlm._get_answer", return_value="Test answer")
llm = CohereLlm(cohere_llm_config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
def test_get_answer_mocked_cohere(cohere_llm_config, mocker):
mocked_cohere = mocker.patch("embedchain.llm.cohere.Cohere")
mock_instance = mocked_cohere.return_value
mock_instance.return_value = "Mocked answer"
llm = CohereLlm(cohere_llm_config)
prompt = "Test query"
answer = llm.get_llm_model_answer(prompt)
assert answer == "Mocked answer"
mocked_cohere.assert_called_once_with(
cohere_api_key="test_api_key",
model="gptd-instruct-tft",
max_tokens=50,
temperature=0.7,
p=0.8,
)
mock_instance.assert_called_once_with(prompt)
+45 -48
View File
@@ -1,64 +1,61 @@
import importlib
import os
import unittest
from unittest.mock import MagicMock, patch
import pytest
from embedchain.config import BaseLlmConfig
from embedchain.llm.huggingface import HuggingFaceLlm
class TestHuggingFaceLlm(unittest.TestCase):
def setUp(self):
os.environ["HUGGINGFACE_ACCESS_TOKEN"] = "test_access_token"
self.config = BaseLlmConfig(model="google/flan-t5-xxl", max_tokens=50, temperature=0.7, top_p=0.8)
@pytest.fixture
def huggingface_llm_config():
os.environ["HUGGINGFACE_ACCESS_TOKEN"] = "test_access_token"
config = BaseLlmConfig(model="google/flan-t5-xxl", max_tokens=50, temperature=0.7, top_p=0.8)
yield config
os.environ.pop("HUGGINGFACE_ACCESS_TOKEN")
def test_init_raises_value_error_without_api_key(self):
os.environ.pop("HUGGINGFACE_ACCESS_TOKEN")
with self.assertRaises(ValueError):
HuggingFaceLlm()
def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
llm = HuggingFaceLlm(self.config)
llm.config.system_prompt = "system_prompt"
with self.assertRaises(ValueError):
llm.get_llm_model_answer("prompt")
def test_init_raises_value_error_without_api_key(mocker):
mocker.patch.dict(os.environ, clear=True)
with pytest.raises(ValueError):
HuggingFaceLlm()
def test_top_p_value_within_range(self):
config = BaseLlmConfig(top_p=1.0)
with self.assertRaises(ValueError):
HuggingFaceLlm._get_answer("test_prompt", config)
def test_dependency_is_imported(self):
importlib_installed = True
try:
importlib.import_module("huggingface_hub")
except ImportError:
importlib_installed = False
self.assertTrue(importlib_installed)
def test_get_llm_model_answer_raises_value_error_for_system_prompt(huggingface_llm_config):
llm = HuggingFaceLlm(huggingface_llm_config)
llm.config.system_prompt = "system_prompt"
with pytest.raises(ValueError):
llm.get_llm_model_answer("prompt")
@patch("embedchain.llm.huggingface.HuggingFaceLlm._get_answer")
def test_get_llm_model_answer(self, mock_get_answer):
mock_get_answer.return_value = "Test answer"
llm = HuggingFaceLlm(self.config)
answer = llm.get_llm_model_answer("Test query")
def test_top_p_value_within_range():
config = BaseLlmConfig(top_p=1.0)
with pytest.raises(ValueError):
HuggingFaceLlm._get_answer("test_prompt", config)
self.assertEqual(answer, "Test answer")
mock_get_answer.assert_called_once()
@patch("embedchain.llm.huggingface.HuggingFaceHub")
def test_hugging_face_mock(self, mock_huggingface):
mock_llm_instance = MagicMock()
mock_llm_instance.return_value = "Test answer"
mock_huggingface.return_value = mock_llm_instance
def test_dependency_is_imported():
importlib_installed = True
try:
importlib.import_module("huggingface_hub")
except ImportError:
importlib_installed = False
assert importlib_installed
llm = HuggingFaceLlm(self.config)
answer = llm.get_llm_model_answer("Test query")
self.assertEqual(answer, "Test answer")
mock_huggingface.assert_called_once_with(
huggingfacehub_api_token="test_access_token",
repo_id="google/flan-t5-xxl",
model_kwargs={"temperature": 0.7, "max_new_tokens": 50, "top_p": 0.8},
)
mock_llm_instance.assert_called_once_with("Test query")
def test_get_llm_model_answer(huggingface_llm_config, mocker):
mocker.patch("embedchain.llm.huggingface.HuggingFaceLlm._get_answer", return_value="Test answer")
llm = HuggingFaceLlm(huggingface_llm_config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
def test_hugging_face_mock(huggingface_llm_config, mocker):
mock_llm_instance = mocker.Mock(return_value="Test answer")
mocker.patch("embedchain.llm.huggingface.HuggingFaceHub", return_value=mock_llm_instance)
llm = HuggingFaceLlm(huggingface_llm_config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mock_llm_instance.assert_called_once_with("Test query")
+64 -28
View File
@@ -1,40 +1,76 @@
import os
import unittest
from unittest.mock import patch
import pytest
from embedchain.config import BaseLlmConfig
from embedchain.llm.jina import JinaLlm
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
class TestJinaLlm(unittest.TestCase):
def setUp(self):
os.environ["JINACHAT_API_KEY"] = "test_api_key"
self.config = BaseLlmConfig(
temperature=0.7, max_tokens=50, top_p=0.8, stream=False, system_prompt="System prompt"
)
@pytest.fixture
def config():
os.environ["JINACHAT_API_KEY"] = "test_api_key"
config = BaseLlmConfig(temperature=0.7, max_tokens=50, top_p=0.8, stream=False, system_prompt="System prompt")
yield config
os.environ.pop("JINACHAT_API_KEY")
def test_init_raises_value_error_without_api_key(self):
os.environ.pop("JINACHAT_API_KEY")
with self.assertRaises(ValueError):
JinaLlm()
@patch("embedchain.llm.jina.JinaLlm._get_answer")
def test_get_llm_model_answer(self, mock_get_answer):
mock_get_answer.return_value = "Test answer"
def test_init_raises_value_error_without_api_key(mocker):
mocker.patch.dict(os.environ, clear=True)
with pytest.raises(ValueError):
JinaLlm()
llm = JinaLlm(self.config)
answer = llm.get_llm_model_answer("Test query")
self.assertEqual(answer, "Test answer")
mock_get_answer.assert_called_once()
def test_get_llm_model_answer(config, mocker):
mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
@patch("embedchain.llm.jina.JinaLlm._get_answer")
def test_get_llm_model_answer_with_system_prompt(self, mock_get_answer):
self.config.system_prompt = "Custom system prompt"
mock_get_answer.return_value = "Test answer"
llm = JinaLlm(config)
answer = llm.get_llm_model_answer("Test query")
llm = JinaLlm(self.config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("Test query", config)
self.assertEqual(answer, "Test answer")
mock_get_answer.assert_called_once()
def test_get_llm_model_answer_with_system_prompt(config, mocker):
config.system_prompt = "Custom system prompt"
mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
llm = JinaLlm(config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("Test query", config)
def test_get_llm_model_answer_empty_prompt(config, mocker):
mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
llm = JinaLlm(config)
answer = llm.get_llm_model_answer("")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("", config)
def test_get_llm_model_answer_with_streaming(config, mocker):
config.stream = True
mocked_jinachat = mocker.patch("embedchain.llm.jina.JinaChat")
llm = JinaLlm(config)
llm.get_llm_model_answer("Test query")
mocked_jinachat.assert_called_once()
callbacks = [callback[1]["callbacks"] for callback in mocked_jinachat.call_args_list]
assert any(isinstance(callback[0], StreamingStdOutCallbackHandler) for callback in callbacks)
def test_get_llm_model_answer_without_system_prompt(config, mocker):
config.system_prompt = None
mocked_jinachat = mocker.patch("embedchain.llm.jina.JinaChat")
llm = JinaLlm(config)
llm.get_llm_model_answer("Test query")
mocked_jinachat.assert_called_once_with(
temperature=config.temperature,
max_tokens=config.max_tokens,
model_kwargs={"top_p": config.top_p},
)
+47
View File
@@ -0,0 +1,47 @@
import os
import pytest
from embedchain.llm.llama2 import Llama2Llm
@pytest.fixture
def llama2_llm():
os.environ["REPLICATE_API_TOKEN"] = "test_api_token"
llm = Llama2Llm()
return llm
def test_init_raises_value_error_without_api_key(mocker):
mocker.patch.dict(os.environ, clear=True)
with pytest.raises(ValueError):
Llama2Llm()
def test_get_llm_model_answer_raises_value_error_for_system_prompt(llama2_llm):
llama2_llm.config.system_prompt = "system_prompt"
with pytest.raises(ValueError):
llama2_llm.get_llm_model_answer("prompt")
def test_get_llm_model_answer(llama2_llm, mocker):
mocked_replicate = mocker.patch("embedchain.llm.llama2.Replicate")
mocked_replicate_instance = mocker.MagicMock()
mocked_replicate.return_value = mocked_replicate_instance
mocked_replicate_instance.return_value = "Test answer"
llama2_llm.config.model = "test_model"
llama2_llm.config.max_tokens = 50
llama2_llm.config.temperature = 0.7
llama2_llm.config.top_p = 0.8
answer = llama2_llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mocked_replicate.assert_called_once_with(
model="test_model",
input={
"temperature": 0.7,
"max_length": 50,
"top_p": 0.8,
},
)
mocked_replicate_instance.assert_called_once_with("Test query")
+73
View File
@@ -0,0 +1,73 @@
import os
import pytest
from embedchain.config import BaseLlmConfig
from embedchain.llm.openai import OpenAILlm
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
@pytest.fixture
def config():
os.environ["OPENAI_API_KEY"] = "test_api_key"
config = BaseLlmConfig(
temperature=0.7, max_tokens=50, top_p=0.8, stream=False, system_prompt="System prompt", model="gpt-3.5-turbo"
)
yield config
os.environ.pop("OPENAI_API_KEY")
def test_get_llm_model_answer(config, mocker):
mocked_get_answer = mocker.patch("embedchain.llm.openai.OpenAILlm._get_answer", return_value="Test answer")
llm = OpenAILlm(config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("Test query", config)
def test_get_llm_model_answer_with_system_prompt(config, mocker):
config.system_prompt = "Custom system prompt"
mocked_get_answer = mocker.patch("embedchain.llm.openai.OpenAILlm._get_answer", return_value="Test answer")
llm = OpenAILlm(config)
answer = llm.get_llm_model_answer("Test query")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("Test query", config)
def test_get_llm_model_answer_empty_prompt(config, mocker):
mocked_get_answer = mocker.patch("embedchain.llm.openai.OpenAILlm._get_answer", return_value="Test answer")
llm = OpenAILlm(config)
answer = llm.get_llm_model_answer("")
assert answer == "Test answer"
mocked_get_answer.assert_called_once_with("", config)
def test_get_llm_model_answer_with_streaming(config, mocker):
config.stream = True
mocked_jinachat = mocker.patch("embedchain.llm.openai.ChatOpenAI")
llm = OpenAILlm(config)
llm.get_llm_model_answer("Test query")
mocked_jinachat.assert_called_once()
callbacks = [callback[1]["callbacks"] for callback in mocked_jinachat.call_args_list]
assert any(isinstance(callback[0], StreamingStdOutCallbackHandler) for callback in callbacks)
def test_get_llm_model_answer_without_system_prompt(config, mocker):
config.system_prompt = None
mocked_jinachat = mocker.patch("embedchain.llm.openai.ChatOpenAI")
llm = OpenAILlm(config)
llm.get_llm_model_answer("Test query")
mocked_jinachat.assert_called_once_with(
model=config.model,
temperature=config.temperature,
max_tokens=config.max_tokens,
model_kwargs={"top_p": config.top_p},
)
+75 -125
View File
@@ -1,135 +1,85 @@
import os
import unittest
import pytest
from unittest.mock import MagicMock, patch
from embedchain import App
from embedchain.config import AppConfig, BaseLlmConfig
class TestApp(unittest.TestCase):
os.environ["OPENAI_API_KEY"] = "test_key"
@pytest.fixture
def app():
os.environ["OPENAI_API_KEY"] = "test_api_key"
app = App(config=AppConfig(collect_metrics=False))
return app
def setUp(self):
self.app = App(config=AppConfig(collect_metrics=False))
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query(self):
"""
This test checks the functionality of the 'query' method in the App class.
It simulates a scenario where the 'retrieve_from_database' method returns a context list and
'get_llm_model_answer' returns an expected answer string.
The 'query' method is expected to call 'retrieve_from_database' and 'get_llm_model_answer' methods
appropriately and return the right answer.
Key assumptions tested:
- 'retrieve_from_database' method is called exactly once with arguments: "Test query" and an instance of
BaseLlmConfig.
- 'get_llm_model_answer' is called exactly once. The specific arguments are not checked in this test.
- 'query' method returns the value it received from 'get_llm_model_answer'.
The test isolates the 'query' method behavior by mocking out 'retrieve_from_database' and
'get_llm_model_answer' methods.
"""
with patch.object(self.app, "retrieve_from_database") as mock_retrieve:
mock_retrieve.return_value = ["Test context"]
with patch.object(self.app.llm, "get_llm_model_answer") as mock_answer:
mock_answer.return_value = "Test answer"
_answer = self.app.query(input_query="Test query")
# Ensure retrieve_from_database was called
mock_retrieve.assert_called_once()
# Check the call arguments
args, kwargs = mock_retrieve.call_args
input_query_arg = kwargs.get("input_query")
self.assertEqual(input_query_arg, "Test query")
mock_answer.assert_called_once()
@patch("embedchain.llm.openai.OpenAILlm._get_answer")
def test_query_config_app_passing(self, mock_get_answer):
mock_get_answer.return_value = MagicMock()
mock_get_answer.return_value.content = "Test answer"
config = AppConfig(collect_metrics=False)
chat_config = BaseLlmConfig(system_prompt="Test system prompt")
app = App(config=config, llm_config=chat_config)
answer = app.llm.get_llm_model_answer("Test query")
self.assertEqual(app.llm.config.system_prompt, "Test system prompt")
self.assertEqual(answer, "Test answer")
@patch("embedchain.llm.openai.OpenAILlm._get_answer")
def test_app_passing(self, mock_get_answer):
mock_get_answer.return_value = MagicMock()
mock_get_answer.return_value.content = "Test answer"
config = AppConfig(collect_metrics=False)
chat_config = BaseLlmConfig()
app = App(config=config, llm_config=chat_config, system_prompt="Test system prompt")
answer = app.llm.get_llm_model_answer("Test query")
self.assertEqual(app.llm.config.system_prompt, "Test system prompt")
self.assertEqual(answer, "Test answer")
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query_with_where_in_params(self):
"""
This test checks the functionality of the 'query' method in the App class.
It simulates a scenario where the 'retrieve_from_database' method returns a context list based on
a where filter and 'get_llm_model_answer' returns an expected answer string.
The 'query' method is expected to call 'retrieve_from_database' with the where filter and
'get_llm_model_answer' methods appropriately and return the right answer.
Key assumptions tested:
- 'retrieve_from_database' method is called exactly once with arguments: "Test query" and an instance of
BaseLlmConfig.
- 'get_llm_model_answer' is called exactly once. The specific arguments are not checked in this test.
- 'query' method returns the value it received from 'get_llm_model_answer'.
The test isolates the 'query' method behavior by mocking out 'retrieve_from_database' and
'get_llm_model_answer' methods.
"""
with patch.object(self.app, "retrieve_from_database") as mock_retrieve:
mock_retrieve.return_value = ["Test context"]
with patch.object(self.app.llm, "get_llm_model_answer") as mock_answer:
mock_answer.return_value = "Test answer"
answer = self.app.query("Test query", where={"attribute": "value"})
self.assertEqual(answer, "Test answer")
_args, kwargs = mock_retrieve.call_args
self.assertEqual(kwargs.get("input_query"), "Test query")
self.assertEqual(kwargs.get("where"), {"attribute": "value"})
mock_answer.assert_called_once()
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query_with_where_in_query_config(self):
"""
This test checks the functionality of the 'query' method in the App class.
It simulates a scenario where the 'retrieve_from_database' method returns a context list based on
a where filter and 'get_llm_model_answer' returns an expected answer string.
The 'query' method is expected to call 'retrieve_from_database' with the where filter and
'get_llm_model_answer' methods appropriately and return the right answer.
Key assumptions tested:
- 'retrieve_from_database' method is called exactly once with arguments: "Test query" and an instance of
BaseLlmConfig.
- 'get_llm_model_answer' is called exactly once. The specific arguments are not checked in this test.
- 'query' method returns the value it received from 'get_llm_model_answer'.
The test isolates the 'query' method behavior by mocking out 'retrieve_from_database' and
'get_llm_model_answer' methods.
"""
with patch.object(self.app.llm, "get_llm_model_answer") as mock_answer:
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query(app):
with patch.object(app, "retrieve_from_database") as mock_retrieve:
mock_retrieve.return_value = ["Test context"]
with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
mock_answer.return_value = "Test answer"
with patch.object(self.app.db, "query") as mock_database_query:
mock_database_query.return_value = ["Test context"]
llm_config = BaseLlmConfig(where={"attribute": "value"})
answer = self.app.query("Test query", llm_config)
answer = app.query(input_query="Test query")
assert answer == "Test answer"
self.assertEqual(answer, "Test answer")
_args, kwargs = mock_database_query.call_args
self.assertEqual(kwargs.get("input_query"), "Test query")
self.assertEqual(kwargs.get("where"), {"attribute": "value"})
mock_answer.assert_called_once()
mock_retrieve.assert_called_once()
_, kwargs = mock_retrieve.call_args
input_query_arg = kwargs.get("input_query")
assert input_query_arg == "Test query"
mock_answer.assert_called_once()
@patch("embedchain.llm.openai.OpenAILlm._get_answer")
def test_query_config_app_passing(mock_get_answer):
mock_get_answer.return_value = MagicMock()
mock_get_answer.return_value = "Test answer"
config = AppConfig(collect_metrics=False)
chat_config = BaseLlmConfig(system_prompt="Test system prompt")
app = App(config=config, llm_config=chat_config)
answer = app.llm.get_llm_model_answer("Test query")
assert app.llm.config.system_prompt == "Test system prompt"
assert answer == "Test answer"
@patch("embedchain.llm.openai.OpenAILlm._get_answer")
def test_app_passing(mock_get_answer):
mock_get_answer.return_value = MagicMock()
mock_get_answer.return_value = "Test answer"
config = AppConfig(collect_metrics=False)
chat_config = BaseLlmConfig()
app = App(config=config, llm_config=chat_config, system_prompt="Test system prompt")
answer = app.llm.get_llm_model_answer("Test query")
assert app.llm.config.system_prompt == "Test system prompt"
assert answer == "Test answer"
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query_with_where_in_params(app):
with patch.object(app, "retrieve_from_database") as mock_retrieve:
mock_retrieve.return_value = ["Test context"]
with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
mock_answer.return_value = "Test answer"
answer = app.query("Test query", where={"attribute": "value"})
assert answer == "Test answer"
_, kwargs = mock_retrieve.call_args
assert kwargs.get("input_query") == "Test query"
assert kwargs.get("where") == {"attribute": "value"}
mock_answer.assert_called_once()
@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
def test_query_with_where_in_query_config(app):
with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
mock_answer.return_value = "Test answer"
with patch.object(app.db, "query") as mock_database_query:
mock_database_query.return_value = ["Test context"]
llm_config = BaseLlmConfig(where={"attribute": "value"})
answer = app.query("Test query", llm_config)
assert answer == "Test answer"
_, kwargs = mock_database_query.call_args
assert kwargs.get("input_query") == "Test query"
assert kwargs.get("where") == {"attribute": "value"}
mock_answer.assert_called_once()
+24 -27
View File
@@ -1,32 +1,29 @@
import unittest
from embedchain.models.data_type import (DataType, DirectDataType,
IndirectDataType, SpecialDataType)
from embedchain.models.data_type import DataType, DirectDataType, IndirectDataType, SpecialDataType
class TestDataTypeEnums(unittest.TestCase):
def test_subclass_types_in_data_type(self):
"""Test that all data type category subclasses are contained in the composite data type"""
# Check if DirectDataType values are in DataType
for data_type in DirectDataType:
self.assertIn(data_type.value, DataType._value2member_map_)
def test_subclass_types_in_data_type():
"""Test that all data type category subclasses are contained in the composite data type"""
# Check if DirectDataType values are in DataType
for data_type in DirectDataType:
assert data_type.value in DataType._value2member_map_
# Check if IndirectDataType values are in DataType
for data_type in IndirectDataType:
self.assertIn(data_type.value, DataType._value2member_map_)
# Check if IndirectDataType values are in DataType
for data_type in IndirectDataType:
assert data_type.value in DataType._value2member_map_
# Check if SpecialDataType values are in DataType
for data_type in SpecialDataType:
self.assertIn(data_type.value, DataType._value2member_map_)
# Check if SpecialDataType values are in DataType
for data_type in SpecialDataType:
assert data_type.value in DataType._value2member_map_
def test_data_type_in_subclasses(self):
"""Test that all data types in the composite data type are categorized in a subclass"""
for data_type in DataType:
if data_type.value in DirectDataType._value2member_map_:
self.assertIn(data_type.value, DirectDataType._value2member_map_)
elif data_type.value in IndirectDataType._value2member_map_:
self.assertIn(data_type.value, IndirectDataType._value2member_map_)
elif data_type.value in SpecialDataType._value2member_map_:
self.assertIn(data_type.value, SpecialDataType._value2member_map_)
else:
self.fail(f"{data_type.value} not found in any subclass enums")
def test_data_type_in_subclasses():
"""Test that all data types in the composite data type are categorized in a subclass"""
for data_type in DataType:
if data_type.value in DirectDataType._value2member_map_:
assert data_type.value in DirectDataType._value2member_map_
elif data_type.value in IndirectDataType._value2member_map_:
assert data_type.value in IndirectDataType._value2member_map_
elif data_type.value in SpecialDataType._value2member_map_:
assert data_type.value in SpecialDataType._value2member_map_
else:
assert False, f"{data_type.value} not found in any subclass enums"
+106
View File
@@ -0,0 +1,106 @@
from unittest import mock
from unittest.mock import patch
from embedchain import App
from embedchain.config import AppConfig
from embedchain.embedder.base import BaseEmbedder
from embedchain.vectordb.pinecone import PineconeDB
class TestPinecone:
@patch("embedchain.vectordb.pinecone.pinecone")
def test_init(self, pinecone_mock):
"""Test that the PineconeDB can be initialized."""
# Create a PineconeDB instance
PineconeDB()
# Assert that the Pinecone client was initialized
pinecone_mock.init.assert_called_once()
pinecone_mock.list_indexes.assert_called_once()
pinecone_mock.Index.assert_called_once()
@patch("embedchain.vectordb.pinecone.pinecone")
def test_set_embedder(self, pinecone_mock):
"""Test that the embedder can be set."""
# Set the embedder
embedder = BaseEmbedder()
# Create a PineconeDB instance
db = PineconeDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
# Assert that the embedder was set
assert db.embedder == embedder
pinecone_mock.init.assert_called_once()
@patch("embedchain.vectordb.pinecone.pinecone")
def test_add_documents(self, pinecone_mock):
"""Test that documents can be added to the database."""
pinecone_client_mock = pinecone_mock.Index.return_value
embedding_function = mock.Mock()
base_embedder = BaseEmbedder()
base_embedder.set_embedding_fn(embedding_function)
vectors = [[0, 0, 0], [1, 1, 1]]
embedding_function.return_value = vectors
# Create a PineconeDb instance
db = PineconeDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=base_embedder)
# Add some documents to the database
documents = ["This is a document.", "This is another document."]
metadatas = [{}, {}]
ids = ["doc1", "doc2"]
db.add(vectors, documents, metadatas, ids, True)
expected_pinecone_upsert_args = [
{"id": "doc1", "metadata": {"text": "This is a document."}, "values": [0, 0, 0]},
{"id": "doc2", "metadata": {"text": "This is another document."}, "values": [1, 1, 1]},
]
# Assert that the Pinecone client was called to upsert the documents
pinecone_client_mock.upsert.assert_called_once_with(expected_pinecone_upsert_args)
@patch("embedchain.vectordb.pinecone.pinecone")
def test_query_documents(self, pinecone_mock):
"""Test that documents can be queried from the database."""
pinecone_client_mock = pinecone_mock.Index.return_value
embedding_function = mock.Mock()
base_embedder = BaseEmbedder()
base_embedder.set_embedding_fn(embedding_function)
vectors = [[0, 0, 0]]
embedding_function.return_value = vectors
# Create a PineconeDB instance
db = PineconeDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=base_embedder)
# Query the database for documents that are similar to "document"
input_query = ["document"]
n_results = 1
db.query(input_query, n_results, where={}, skip_embedding=False)
# Assert that the Pinecone client was called to query the database
pinecone_client_mock.query.assert_called_once_with(
vector=db.embedder.embedding_fn(input_query)[0], top_k=n_results, filter={}, include_metadata=True
)
@patch("embedchain.vectordb.pinecone.pinecone")
def test_reset(self, pinecone_mock):
"""Test that the database can be reset."""
# Create a PineconeDb instance
db = PineconeDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=BaseEmbedder())
# Reset the database
db.reset()
# Assert that the Pinecone client was called to delete the index
pinecone_mock.delete_index.assert_called_once_with(db.index_name)
# Assert that the index is recreated
pinecone_mock.Index.assert_called_with(db.index_name)
+244
View File
@@ -0,0 +1,244 @@
import unittest
from unittest.mock import patch
from embedchain import App
from embedchain.config import AppConfig
from embedchain.config.vectordb.pinecone import PineconeDBConfig
from embedchain.embedder.base import BaseEmbedder
from embedchain.vectordb.weaviate import WeaviateDB
class TestWeaviateDb(unittest.TestCase):
def test_incorrect_config_throws_error(self):
"""Test the init method of the WeaviateDb class throws error for incorrect config"""
with self.assertRaises(TypeError):
WeaviateDB(config=PineconeDBConfig())
@patch("embedchain.vectordb.weaviate.weaviate")
def test_initialize(self, weaviate_mock):
"""Test the init method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_schema_mock = weaviate_client_mock.schema
# Mock that schema doesn't already exist so that a new schema is created
weaviate_client_schema_mock.exists.return_value = False
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
expected_class_obj = {
"classes": [
{
"class": "Embedchain_store_1526",
"vectorizer": "none",
"properties": [
{
"name": "identifier",
"dataType": ["text"],
},
{
"name": "text",
"dataType": ["text"],
},
{
"name": "metadata",
"dataType": ["Embedchain_store_1526_metadata"],
},
],
},
{
"class": "Embedchain_store_1526_metadata",
"vectorizer": "none",
"properties": [
{
"name": "data_type",
"dataType": ["text"],
},
{
"name": "doc_id",
"dataType": ["text"],
},
{
"name": "url",
"dataType": ["text"],
},
{
"name": "hash",
"dataType": ["text"],
},
{
"name": "app_id",
"dataType": ["text"],
},
{
"name": "text",
"dataType": ["text"],
},
],
},
]
}
# Assert that the Weaviate client was initialized
weaviate_mock.Client.assert_called_once()
self.assertEqual(db.index_name, "Embedchain_store_1526")
weaviate_client_schema_mock.create.assert_called_once_with(expected_class_obj)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_get_or_create_db(self, weaviate_mock):
"""Test the _get_or_create_db method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
expected_client = db._get_or_create_db()
self.assertEqual(expected_client, weaviate_client_mock)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_add(self, weaviate_mock):
"""Test the add method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_batch_mock = weaviate_client_mock.batch
weaviate_client_batch_enter_mock = weaviate_client_mock.batch.__enter__.return_value
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
db.BATCH_SIZE = 1
embeddings = [[1, 2, 3], [4, 5, 6]]
documents = ["This is a test document.", "This is another test document."]
metadatas = [None, None]
ids = ["123", "456"]
skip_embedding = True
db.add(embeddings, documents, metadatas, ids, skip_embedding)
# Check if the document was added to the database.
weaviate_client_batch_mock.configure.assert_called_once_with(batch_size=1, timeout_retries=3)
weaviate_client_batch_enter_mock.add_data_object.assert_any_call(
data_object={"text": documents[0]}, class_name="Embedchain_store_1526_metadata", vector=embeddings[0]
)
weaviate_client_batch_enter_mock.add_data_object.assert_any_call(
data_object={"text": documents[1]}, class_name="Embedchain_store_1526_metadata", vector=embeddings[1]
)
weaviate_client_batch_enter_mock.add_data_object.assert_any_call(
data_object={"identifier": ids[0], "text": documents[0]},
class_name="Embedchain_store_1526",
vector=embeddings[0],
)
weaviate_client_batch_enter_mock.add_data_object.assert_any_call(
data_object={"identifier": ids[1], "text": documents[1]},
class_name="Embedchain_store_1526",
vector=embeddings[1],
)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_query_without_where(self, weaviate_mock):
"""Test the query method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_query_mock = weaviate_client_mock.query
weaviate_client_query_get_mock = weaviate_client_query_mock.get.return_value
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
# Query for the document.
db.query(input_query=["This is a test document."], n_results=1, where={}, skip_embedding=True)
weaviate_client_query_mock.get.assert_called_once_with("Embedchain_store_1526", ["text"])
weaviate_client_query_get_mock.with_near_vector.assert_called_once_with(
{"vector": ["This is a test document."]}
)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_query_with_where(self, weaviate_mock):
"""Test the query method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_query_mock = weaviate_client_mock.query
weaviate_client_query_get_mock = weaviate_client_query_mock.get.return_value
weaviate_client_query_get_where_mock = weaviate_client_query_get_mock.with_where.return_value
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
# Query for the document.
db.query(input_query=["This is a test document."], n_results=1, where={"doc_id": "123"}, skip_embedding=True)
weaviate_client_query_mock.get.assert_called_once_with("Embedchain_store_1526", ["text"])
weaviate_client_query_get_mock.with_where.assert_called_once_with(
{"operator": "Equal", "path": ["metadata", "Embedchain_store_1526_metadata", "doc_id"], "valueText": "123"}
)
weaviate_client_query_get_where_mock.with_near_vector.assert_called_once_with(
{"vector": ["This is a test document."]}
)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_reset(self, weaviate_mock):
"""Test the reset method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_batch_mock = weaviate_client_mock.batch
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
# Reset the database.
db.reset()
weaviate_client_batch_mock.delete_objects.assert_called_once_with(
"Embedchain_store_1526", where={"path": ["identifier"], "operator": "Like", "valueText": ".*"}
)
@patch("embedchain.vectordb.weaviate.weaviate")
def test_count(self, weaviate_mock):
"""Test the reset method of the WeaviateDb class."""
weaviate_client_mock = weaviate_mock.Client.return_value
weaviate_client_query = weaviate_client_mock.query
# Set the embedder
embedder = BaseEmbedder()
embedder.set_vector_dimension(1526)
# Create a Weaviate instance
db = WeaviateDB()
app_config = AppConfig(collect_metrics=False)
App(config=app_config, db=db, embedder=embedder)
# Reset the database.
db.count()
weaviate_client_query.aggregate.assert_called_once_with("Embedchain_store_1526")