Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bb490df9a6 | ||
|
|
cdfd6519c8 | ||
|
|
e8a2846449 | ||
|
|
d065cbf934 | ||
|
|
413b107b9a | ||
|
|
c336292346 | ||
|
|
adf50f1e81 | ||
|
|
636bc0a99d | ||
|
|
a7a61fae1d | ||
|
|
5ec12212e4 | ||
|
|
77c90a308e |
+4
-1
@@ -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/
|
||||
@@ -8,7 +8,7 @@
|
||||
[](https://colab.research.google.com/drive/138lMWhENGeEu7Q1-6lNbNTHGLZXBBz_B?usp=sharing)
|
||||
[](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},
|
||||
|
||||
@@ -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
|
||||
@@ -25,7 +25,7 @@ llm:
|
||||
vectordb:
|
||||
provider: chroma
|
||||
config:
|
||||
collection_name: 'full-stack-app'
|
||||
collection_name: 'my-collection-name'
|
||||
dir: db
|
||||
allow_reset: true
|
||||
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
vectordb:
|
||||
provider: pinecone
|
||||
config:
|
||||
metric: cosine
|
||||
vector_dimension: 1536
|
||||
collection_name: my-pinecone-index
|
||||
@@ -0,0 +1,4 @@
|
||||
vectordb:
|
||||
provider: weaviate
|
||||
config:
|
||||
collection_name: my_weaviate_index
|
||||
@@ -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|
|
||||
@@ -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
|
||||
'''
|
||||
```
|
||||
@@ -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?")
|
||||
```
|
||||
@@ -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!
|
||||
|
||||
@@ -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
@@ -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'
|
||||
|
||||
@@ -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" />
|
||||
|
||||
@@ -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,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?
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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("_", "-")
|
||||
@@ -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>
|
||||
|
||||
@@ -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
@@ -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
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
@@ -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},
|
||||
)
|
||||
|
||||
@@ -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")
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user