├── .env.example
├── .gitignore
├── LICENSE
├── README.md
├── app
├── __init__.py
├── app.py
├── static
│ ├── dashboard.js
│ └── styles.css
└── templates
│ ├── dashboard.html
│ └── index.html
├── integrations
├── __init__.py
└── email
│ ├── __init__.py
│ ├── fetcher.py
│ └── gmail.py
├── main.py
├── poetry.lock
├── public
├── client_id_secret.png
├── dashboard.png
├── gmail_api_circled.png
├── google_api_library.png
├── google_api_oauth.png
├── google_cloud_api.png
├── google_credentials_oauth.png
├── google_oauth_json_download.png
├── google_uris.png
└── test_user.png
├── pyproject.toml
├── tasks
├── __init__.py
├── agents.py
├── execution.py
├── processor.py
└── storage.py
├── tests
├── babyagi.py
├── colorlogs.py
├── embedding.py
├── entity_add.py
├── graph_agent.py
├── ollama_raw.py
└── ollama_streaming.py
└── utils
├── __init__.py
├── custom_log_formatter.py
└── ollama.py
/.env.example:
--------------------------------------------------------------------------------
1 | GOOGLE_CLIENT_ID="YOUR_GOOGLE_CLIENT_ID"
2 | GOOGLE_CLIENT_SECRET="YOUR_GOOGLE_CLIENT_SECRET"
3 |
4 | GOOGLE_REDIRECT_URI=http://localhost:8080/oauth2callback
5 | GOOGLE_LOGIN_URI=http://localhost:8080/authorize
6 |
7 | FLASK_SECRET_KEY="YOUR_FLASK_SECRET_KEY"
8 |
9 | NEXUSDB_API_KEY="YOUR_NEXUSDB_API_KEY"
10 |
11 | MAX_THREADS=4
12 | INITIAL_EMAILS=1
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | # C extensions
7 | *.so
8 |
9 | # Distribution / packaging
10 | .Python
11 | build/
12 | develop-eggs/
13 | dist/
14 | downloads/
15 | eggs/
16 | .eggs/
17 | lib/
18 | lib64/
19 | parts/
20 | sdist/
21 | var/
22 | wheels/
23 | share/python-wheels/
24 | *.egg-info/
25 | .installed.cfg
26 | *.egg
27 | MANIFEST
28 |
29 | # PyInstaller
30 | # Usually these files are written by a python script from a template
31 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
32 | *.manifest
33 | *.spec
34 |
35 | # Installer logs
36 | pip-log.txt
37 | pip-delete-this-directory.txt
38 |
39 | # Unit test / coverage reports
40 | htmlcov/
41 | .tox/
42 | .nox/
43 | .coverage
44 | .coverage.*
45 | .cache
46 | nosetests.xml
47 | coverage.xml
48 | *.cover
49 | *.py,cover
50 | .hypothesis/
51 | .pytest_cache/
52 | cover/
53 |
54 | # Translations
55 | *.mo
56 | *.pot
57 |
58 | # Django stuff:
59 | *.log
60 | local_settings.py
61 | db.sqlite3
62 | db.sqlite3-journal
63 |
64 | # Flask stuff:
65 | instance/
66 | .webassets-cache
67 |
68 | # Scrapy stuff:
69 | .scrapy
70 |
71 | # Sphinx documentation
72 | docs/_build/
73 |
74 | # PyBuilder
75 | .pybuilder/
76 | target/
77 |
78 | # Jupyter Notebook
79 | .ipynb_checkpoints
80 |
81 | # IPython
82 | profile_default/
83 | ipython_config.py
84 |
85 | # pyenv
86 | # For a library or package, you might want to ignore these files since the code is
87 | # intended to run in multiple environments; otherwise, check them in:
88 | # .python-version
89 |
90 | # pipenv
91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
94 | # install all needed dependencies.
95 | #Pipfile.lock
96 |
97 | # poetry
98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
99 | # This is especially recommended for binary packages to ensure reproducibility, and is more
100 | # commonly ignored for libraries.
101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102 | #poetry.lock
103 |
104 | # pdm
105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106 | #pdm.lock
107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108 | # in version control.
109 | # https://pdm.fming.dev/#use-with-ide
110 | .pdm.toml
111 |
112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113 | __pypackages__/
114 |
115 | # Celery stuff
116 | celerybeat-schedule
117 | celerybeat.pid
118 |
119 | # SageMath parsed files
120 | *.sage.py
121 |
122 | # Environments
123 | .env
124 | .venv
125 | env/
126 | venv/
127 | ENV/
128 | env.bak/
129 | venv.bak/
130 | poetry.lock
131 |
132 | # Spyder project settings
133 | .spyderproject
134 | .spyproject
135 |
136 | # Rope project settings
137 | .ropeproject
138 |
139 | # mkdocs documentation
140 | /site
141 |
142 | # mypy
143 | .mypy_cache/
144 | .dmypy.json
145 | dmypy.json
146 |
147 | # Pyre type checker
148 | .pyre/
149 |
150 | # pytype static type analyzer
151 | .pytype/
152 |
153 | # Cython debug symbols
154 | cython_debug/
155 |
156 | # VS Code
157 | .vscode/
158 |
159 | # Other
160 | other/
161 | client_secret.json
162 |
163 | # PyCharm
164 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
165 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
166 | # and can be added to the global gitignore or merged into this file. For a more nuclear
167 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
168 | #.idea/
169 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Non-Compete Open License
2 |
3 | Copyright (c) 2024 Astra Analytics, Inc.
4 |
5 | Permission is hereby granted, free of charge, to any person obtaining a copy
6 | of this software and associated documentation files (the "Software"), to deal
7 | in the Software, including without limitation the rights to use, copy, modify,
8 | merge, publish, distribute, sublicense, and/or sell copies of the Software,
9 | and to permit persons to whom the Software is furnished to do so, subject to
10 | the following conditions:
11 |
12 | **Non-Compete Restriction**
13 | You may not modify, change, or replace the database system (NexusDB) that the Software
14 | uses. The Software is designed to work with NexusDB, and any alteration to the
15 | database system is prohibited.
16 |
17 | The above copyright notice and this permission notice shall be included in all
18 | copies or substantial portions of the Software.
19 |
20 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26 | SOFTWARE.
27 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Task Agent Starter Kit with NexusDB
2 |
3 | 
4 |
5 | ## Table of Contents
6 |
7 | - [Task Agent Starter Kit with NexusDB](#task-agent-starter-kit-with-nexusdb)
8 | - [Table of Contents](#table-of-contents)
9 | - [Setup](#setup)
10 | - [Connecting to Gmail](#connecting-to-gmail)
11 | - [Step 1: Create a Project and Enable the Gmail API](#step-1-create-a-project-and-enable-the-gmail-api)
12 | - [Step 2: Configure OAuth Consent Screen](#step-2-configure-oauth-consent-screen)
13 | - [Step 3: Create Credentials](#step-3-create-credentials)
14 | - [Configure The AI Models](#configure-the-ai-models)
15 | - [Step 1: Install Ollama](#step-1-install-ollama)
16 | - [Environment Variables](#environment-variables)
17 | - [NexusDB API Key](#nexusdb-api-key)
18 | - [Max Threads](#max-threads)
19 | - [Initial Emails](#initial-emails)
20 | - [Installation](#installation)
21 | - [Running the app](#running-the-app)
22 | - [Using Poetry Shell](#using-poetry-shell)
23 | - [Using Poetry Run](#using-poetry-run)
24 | - [Contributing](#contributing)
25 | - [License](#license)
26 | - [Non-Compete Open License](#non-compete-open-license)
27 | - [Permissions](#permissions)
28 | - [Conditions](#conditions)
29 | - [Summary](#summary)
30 | - [NexusDB](#nexusdb)
31 | - [Earn 30% Commission on referrals](#earn-30-commission-on-referrals)
32 | - [Partner with us](#partner-with-us)
33 |
34 | ## Setup
35 |
36 | To run this application there are a few items that must be set up separately:
37 |
38 | - Gmail: For the app to connect to your Gmail account securely, you should create your own API credentials. We'll walk through the steps to create an application in Google Cloud and enable the Gmail API.
39 | - AI Models: This application uses [Llama 3](https://ai.meta.com/blog/meta-llama-3/) for the AI agents and [mxbai-embed-large](https://www.mixedbread.ai/docs//mxbai-embed-large-v1#model-description), which are not included as part of the application and must be downloaded separately. Like with Gmail, the steps for setting this up are below.
40 |
41 | ### Connecting to Gmail
42 |
43 | #### Step 1: Create a Project and Enable the Gmail API
44 |
45 | 1. Navigate to [Google Cloud Console](https://console.cloud.google.com/).
46 | 2. Create a new project if you don't already have one.
47 | 3. Enable the Gmail API: In the navigation menu, select “APIs & Services” -> “Library”
48 |
49 | It might not look like this, you may have to search for it in the search bar at the top.
50 |
51 | 
52 |
53 | 
54 |
55 | 4. Search for “Gmail API” and enable it for your project.
56 |
57 | 
58 |
59 | #### Step 2: Configure OAuth Consent Screen
60 |
61 | 1. In the Google Cloud Console, go to “OAuth consent screen” under “APIs & Services”.
62 |
63 | 
64 |
65 | 2. Set the User Type to “External” and create.
66 | 3. Fill in the required fields like App name, User support email, and Developer contact information.
67 | 4. Save and continue until you finish the setup.
68 |
69 | _Note: be sure to add yourself as a test user_
70 | 
71 |
72 | #### Step 3: Create Credentials
73 |
74 | 1. In the Google Cloud Console, under “Credentials” (still within “APIs & Services”), click “Create Credentials” and choose “OAuth Client ID”.
75 |
76 | 
77 |
78 | 2. Select “Web application” as the Application type and give it a name.
79 | 3. Add `http://localhost` and `http://localhost:8080/` as authorized JavaScript origins
80 |
81 | _Note, Google only allows localhost for testing, which is why main.py sets this as host instead of default Flask 127.0.0.1. If you want to use a different port, be sure to make the change here as well as main.py_
82 |
83 | 4. Add `http://localhost:8080/oauth2callback` as authorized redirect URI.
84 |
85 | Your setup should look like this:
86 |
87 | 
88 |
89 | 5. After creating the credentials, click the download button to get your credentials:
90 |
91 | 
92 |
93 | 
94 |
95 | 6. Copy and paste your Google Client ID and Client Secret into the .env file at the root of your project.
96 | 7. **If you don't rename the file from `.env.example` to `.env` you will get an error! Don't miss this step!**
97 |
98 | ### Configure The AI Models
99 |
100 | #### Step 1: Install Ollama
101 |
102 | 1. Go to [Ollama website](https://ollama.com/) and download the application.
103 | 2. After installing, run the following commands in your terminal:
104 |
105 | ```bash
106 | ollama pull llama3
107 | ollama pull mxbai-embed-large
108 | ```
109 |
110 | More on each model here:
111 |
112 | - [Llama 3](https://ai.meta.com/blog/meta-llama-3/)
113 | - [mxbai-embed-large](https://www.mixedbread.ai/docs//mxbai-embed-large-v1#model-description)
114 |
115 | 3. Optionally, you can test the installation by running:
116 |
117 | ```bash
118 | ollama run llama3
119 | ```
120 |
121 | 4. If you want to run more than 1 agent in parallel (recommended), you'll have to configure the ollama server to do so. Make sure the ollama application is not running, then start it yourself in a terminal window using the command:
122 |
123 | ```bash
124 | OLLAMA_NUM_PARALLEL=4 ollama serve
125 | ```
126 |
127 | this will start the ollama server with the number of parallel processes = 4. You can choose any number you'd like, but it's recommended to set this number to double that of the MAX_THREADS environment variable, which we'll discuss in more detail later.
128 |
129 | ### Environment Variables
130 |
131 | **Make sure you have renamed the file from `.env.example` to `.env` or you will get an error!**
132 |
133 | #### NexusDB API Key
134 |
135 | This app runs on NexusDB, so if you don't have an API key yet, go to [nexusdb.io](https://www.nexusdb.io) and sign up for an account. After signing up you will be able to get your API key from the dashboard and paste it into .env
136 |
137 | #### Max Threads
138 |
139 | The MAX_THREADS variable determines the number of emails that can be processed simultaneously. To allow the task agent and graph creation agents to run concurrently for each email, you should set the ollama server to run twice this number of parallel processes
140 |
141 | #### Initial Emails
142 |
143 | This variable sets the number of emails in the inbox the application should add to the queue before waiting for new ones to come in.
144 |
145 | ## Installation
146 |
147 | 1. If you don't have Poetry installed, do that first:
148 |
149 | - **Install Poetry**:
150 | Poetry provides an easy way to install itself. Run the following command:
151 |
152 | ```bash
153 | curl -sSL https://install.python-poetry.org | python3 -
154 | ```
155 |
156 | Alternatively, you can follow the instructions on the [Poetry documentation](https://python-poetry.org/docs/)
157 |
158 | 2. Clone the repository
159 |
160 | 3. Install Dependencies
161 |
162 | Navigate to the project directory and install project dependencies
163 |
164 | ```bash
165 | poetry install
166 | ```
167 |
168 | ## Running the app
169 |
170 | To run the project, activate the Poetry shell or use Poetry's `run` command.
171 |
172 | ### Using Poetry Shell
173 |
174 | Activate the Poetry shell:
175 |
176 | ```bash
177 | poetry shell
178 | ```
179 |
180 | Then, run the project in development mode:
181 |
182 | ```bash
183 | python main.py
184 | ```
185 |
186 | ### Using Poetry Run
187 |
188 | Alternatively, you can use the poetry run command to execute scripts without activating the shell:
189 |
190 | ```bash
191 | poetry run python main.py
192 | ```
193 |
194 | ## Contributing
195 |
196 | We welcome contributions! Please follow these steps to contribute to the project:
197 |
198 | 1. Fork the repository.
199 | 2. Create a new branch (git checkout -b feature-branch).
200 | 3. Make your changes and commit them (git commit -m 'Add some feature').
201 | 4. Push to the branch (git push origin feature-branch).
202 | 5. Open a pull request.
203 |
204 | ## License
205 |
206 | ### Non-Compete Open License
207 |
208 | **Non-Compete Restriction**: This software is designed to work exclusively with NexusDB. You are not permitted to modify, change, or replace the database system used by the software.
209 |
210 | ### Permissions
211 |
212 | You are granted the following rights, free of charge:
213 |
214 | - **Use**: You can use the software for any purpose.
215 | - **Copy**: You can make copies of the software.
216 | - **Modify**: You can modify the software as long as the database system remains NexusDB.
217 | - **Merge**: You can merge the software with other projects.
218 | - **Publish**: You can publish the software.
219 | - **Distribute**: You can distribute the software.
220 | - **Sublicense**: You can sublicense the software.
221 |
222 | ### Conditions
223 |
224 | - Any distribution of the software must include the original copyright notice and this permission notice.
225 | - The software is provided "as-is" without any warranty, express or implied. The authors are not liable for any damages or claims arising from the use of the software.
226 |
227 | ### Summary
228 |
229 | This license grants you broad rights to use, modify, and distribute the software, with the specific condition that the underlying database system must remain NexusDB. This ensures the software's core functionality remains intact and aligned with the intended database system.
230 |
231 | For the full license text, please see the LICENSE file included with this project.
232 |
233 | ## NexusDB
234 |
235 | We on the NexusDB team are building the world's first `data web`, connecting public and private information in a way that's secure, fast, flexible, and highly querable. It's also 20x faster to set up and 50% cheaper for comparable use cases than leading competitors!
236 |
237 | Whether you need to store tables, graphs, embeddings, json objects, or blobs, NexusDB is the solution. See [documentation](https://docs.nexusdb.io) for more details.
238 |
239 | ### Earn 30% Commission on referrals
240 |
241 | If you like NexusDB and want to spread the love, you can get paid to do so through our [referral program](https://www.nexusdb.io/affiliates)! See the website for terms and additional details.
242 |
243 | ### Partner with us
244 |
245 | We're growing fast and are actively seeking design partners and investors! If you enjoyed this demo, we'd love to work with you and tell you about our plans. Reach out to CEO Will Humble at [w@astraanalytics.co](mailto:w@astraanalytics.co) for more info.
246 |
--------------------------------------------------------------------------------
/app/__init__.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 |
4 | from flask import Flask
5 |
6 | from utils.custom_log_formatter import ThreadNameColoredFormatter
7 |
8 | from .app import main
9 |
10 | # Configure colorlog with the custom formatter
11 | formatter = ThreadNameColoredFormatter(
12 | "%(log_color)s[%(threadName)s] - %(message)s",
13 | )
14 |
15 | handler = logging.StreamHandler()
16 | handler.setFormatter(formatter)
17 |
18 | # Suppress logging from specific third-party libraries
19 | logging.getLogger("google_auth_httplib2").setLevel(logging.WARNING)
20 | logging.getLogger("googleapiclient.discovery").setLevel(logging.WARNING)
21 | logging.getLogger("googleapiclient.discovery_cache").setLevel(logging.WARNING)
22 | logging.getLogger("httpcore.http11").setLevel(logging.WARNING)
23 | logging.getLogger("urllib3.connectionpool").setLevel(logging.WARNING)
24 | logging.getLogger("httpcore.connection").setLevel(logging.WARNING)
25 | logging.getLogger("werkzeug").setLevel(logging.WARNING)
26 |
27 | logger = logging.getLogger()
28 | logger.addHandler(handler)
29 | logger.setLevel(logging.INFO)
30 |
31 |
32 | def create_app():
33 | app = Flask(__name__, template_folder="./templates", static_folder="./static")
34 | app.secret_key = os.getenv("FLASK_SECRET_KEY")
35 |
36 | app.register_blueprint(main)
37 |
38 | return app
39 |
--------------------------------------------------------------------------------
/app/app.py:
--------------------------------------------------------------------------------
1 | import os
2 | import threading
3 | from calendar import c
4 |
5 | import flask
6 | import google.oauth2.credentials
7 | import google_auth_oauthlib.flow
8 | import requests
9 | from flask import Blueprint, g, jsonify, redirect, render_template, session, url_for
10 |
11 | from integrations.email.fetcher import email_fetcher
12 | from tasks.processor import start_processing, tasks_storage
13 |
14 | CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
15 | CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
16 | REDIRECT_URI = os.getenv("GOOGLE_REDIRECT_URI")
17 | GOOGLE_LOGIN_URI = os.getenv("GOOGLE_LOGIN_URI")
18 | SCOPES = [
19 | "https://www.googleapis.com/auth/gmail.readonly",
20 | "openid",
21 | "https://www.googleapis.com/auth/userinfo.profile",
22 | "https://www.googleapis.com/auth/userinfo.email",
23 | ]
24 |
25 | main = Blueprint("main", __name__)
26 |
27 |
28 | @main.route("/")
29 | def index():
30 | return render_template("index.html", google_client_id=CLIENT_ID)
31 |
32 |
33 | @main.route("/dashboard", methods=["GET", "POST"])
34 | def dashboard():
35 | tasks = tasks_storage.get_tasks(condition="actionStatus = 'Active'")
36 | agent_tasks = [task for task in tasks.values() if task["agent"] == "AI"]
37 | human_tasks = [task for task in tasks.values() if task["agent"] != "AI"]
38 |
39 | if "credentials" not in session:
40 | return redirect(url_for("main.index"))
41 |
42 | # Start the background processes for email fetching and task processing
43 | if not hasattr(g, "email_fetcher_thread"):
44 | start_processing()
45 | email_fetcher_thread = threading.Thread(
46 | target=email_fetcher,
47 | args=(session["credentials"],),
48 | daemon=True,
49 | name="email_fetcher",
50 | )
51 |
52 | email_fetcher_thread.start()
53 | g.email_fetcher_thread = email_fetcher_thread
54 |
55 | return render_template(
56 | "dashboard.html", agent_tasks=agent_tasks, human_tasks=human_tasks
57 | )
58 |
59 |
60 | @main.route("/tasks")
61 | def get_tasks():
62 | tasks = tasks_storage.get_tasks(condition="actionStatus = 'Active'")
63 | agent_tasks = [task for task in tasks.values() if task["agent"] == "AI"]
64 | human_tasks = [task for task in tasks.values() if task["agent"] != "AI"]
65 | return jsonify(agent_tasks=agent_tasks, human_tasks=human_tasks)
66 |
67 |
68 | @main.route("/authorize", methods=["GET", "POST"])
69 | def authorize():
70 | flow = google_auth_oauthlib.flow.Flow.from_client_config(
71 | {
72 | "web": {
73 | "client_id": CLIENT_ID,
74 | "client_secret": CLIENT_SECRET,
75 | "redirect_uris": [REDIRECT_URI],
76 | "auth_uri": "https://accounts.google.com/o/oauth2/auth",
77 | "token_uri": "https://oauth2.googleapis.com/token",
78 | }
79 | },
80 | scopes=SCOPES,
81 | )
82 | flow.redirect_uri = url_for("main.oauth2callback", _external=True)
83 | authorization_url, state = flow.authorization_url(
84 | access_type="offline", include_granted_scopes="true"
85 | )
86 | session["state"] = state
87 | return redirect(authorization_url)
88 |
89 |
90 | @main.route("/oauth2callback", methods=["GET", "POST"])
91 | def oauth2callback():
92 | state = session["state"]
93 | flow = google_auth_oauthlib.flow.Flow.from_client_config(
94 | {
95 | "web": {
96 | "client_id": CLIENT_ID,
97 | "client_secret": CLIENT_SECRET,
98 | "redirect_uris": [REDIRECT_URI],
99 | "auth_uri": "https://accounts.google.com/o/oauth2/auth",
100 | "token_uri": "https://oauth2.googleapis.com/token",
101 | }
102 | },
103 | scopes=SCOPES,
104 | state=state,
105 | )
106 | flow.redirect_uri = url_for("main.oauth2callback", _external=True)
107 | authorization_response = flask.request.url
108 | flow.fetch_token(authorization_response=authorization_response)
109 |
110 | credentials = flow.credentials
111 | session["credentials"] = credentials_to_dict(credentials)
112 |
113 | return redirect(url_for("main.dashboard"))
114 |
115 |
116 | @main.route("/revoke")
117 | def revoke():
118 | if "credentials" not in session:
119 | return 'You need to authorize before testing the code to revoke credentials.'
120 |
121 | credentials = google.oauth2.credentials.Credentials(**session["credentials"])
122 |
123 | revoke = requests.post(
124 | "https://oauth2.googleapis.com/revoke",
125 | params={"token": credentials.token},
126 | headers={"content-type": "application/x-www-form-urlencoded"},
127 | )
128 |
129 | status_code = getattr(revoke, "status_code")
130 | if status_code == 200:
131 | return "Credentials successfully revoked."
132 | else:
133 | return "An error occurred."
134 |
135 |
136 | @main.route("/clear")
137 | def clear_credentials():
138 | if "credentials" in session:
139 | del session["credentials"]
140 | return "Credentials have been cleared."
141 |
142 |
143 | def credentials_to_dict(credentials):
144 | return {
145 | "token": credentials.token,
146 | "refresh_token": credentials.refresh_token,
147 | "token_uri": credentials.token_uri,
148 | "client_id": credentials.client_id,
149 | "client_secret": credentials.client_secret,
150 | "scopes": credentials.scopes,
151 | }
152 |
--------------------------------------------------------------------------------
/app/static/dashboard.js:
--------------------------------------------------------------------------------
1 | document.addEventListener("DOMContentLoaded", function () {
2 | function fetchTasks() {
3 | console.log("Fetching tasks...");
4 | fetch("/tasks")
5 | .then((response) => {
6 | console.log("Response received:", response);
7 | return response.json();
8 | })
9 | .then((data) => {
10 | console.log("Data received:", data);
11 | updateTasks(data.agent_tasks, "agent-tasks");
12 | updateTasks(data.human_tasks, "human-tasks");
13 | })
14 | .catch((error) => console.error("Error fetching tasks:", error));
15 | }
16 |
17 | function updateTasks(tasks, elementId) {
18 | const tasksList = document.getElementById(elementId);
19 | tasksList.innerHTML = "";
20 |
21 | tasks.forEach((task) => {
22 | const taskItem = document.createElement("li");
23 | taskItem.textContent = task.name;
24 |
25 | tasksList.appendChild(taskItem);
26 |
27 | if (task.potentialAction) {
28 | const subTasksList = document.createElement("ul");
29 | task.potentialAction.forEach((subtask) => {
30 | const subtaskItem = document.createElement("li");
31 | subtaskItem.textContent = subtask;
32 | subTasksList.appendChild(subtaskItem);
33 | });
34 | tasksList.appendChild(subTasksList);
35 | }
36 | });
37 | }
38 |
39 | // Fetch tasks every 10 seconds
40 | setInterval(fetchTasks, 5000);
41 | fetchTasks();
42 | });
43 |
--------------------------------------------------------------------------------
/app/static/styles.css:
--------------------------------------------------------------------------------
1 | /* Basic reset */
2 | * {
3 | margin: 0;
4 | padding: 0;
5 | box-sizing: border-box;
6 | }
7 |
8 | /* Body styling */
9 | body {
10 | font-family: Arial, sans-serif;
11 | background-color: #f4f4f9;
12 | color: #333;
13 | line-height: 1.6;
14 | padding: 20px;
15 | }
16 |
17 | /* Main title */
18 | h1 {
19 | text-align: center;
20 | margin-bottom: 20px;
21 | }
22 |
23 | /* Section titles */
24 | h2 {
25 | color: #555;
26 | margin-bottom: 10px;
27 | }
28 |
29 | /* Task lists */
30 | ul {
31 | list-style-type: none;
32 | margin-bottom: 20px;
33 | }
34 |
35 | li {
36 | background: #fff;
37 | margin: 5px 0;
38 | padding: 10px;
39 | border-radius: 5px;
40 | box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
41 | }
42 |
43 | /* Subtask lists */
44 | ul ul {
45 | margin-left: 20px;
46 | }
47 |
48 | ul ul li {
49 | background: #f9f9f9;
50 | margin: 3px 0;
51 | padding: 8px;
52 | border-radius: 3px;
53 | }
54 |
55 | /* Links */
56 | a {
57 | color: #333;
58 | text-decoration: none;
59 | }
60 |
61 | a:hover {
62 | text-decoration: underline;
63 | }
64 |
--------------------------------------------------------------------------------
/app/templates/dashboard.html:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | Dashboard
6 |
10 |
11 |
12 |
Task Dashboard
13 |
Agent Tasks
14 |
15 | {% for task in agent_tasks %}
16 |
{{ task['name'] }}
17 | {% if task['potentialAction'] %}
18 |
19 | {% for subtask in task['potentialAction'] %}
20 |
{{ subtask }}
21 | {% endfor %}
22 |
23 | {% endif %} {% endfor %}
24 |
25 |
26 |
Human Tasks
27 |
28 | {% for task in human_tasks %}
29 |
{{ task['name'] }}
30 | {% if task['potentialAction'] %}
31 |
32 | {% for subtask in task['potentialAction'] %}
33 |
Please write a haiku about golf. If you have any questions, please email me at john.doe@example.com.
Thank you, John Doe
\r\n\r\n',
17 | "From": "John Doe ",
18 | "Message-ID": "01j0r9h83tfjyrjj8trk7zxe6v",
19 | "Subject": "Haiku about golf",
20 | "Timestamp": "Tue, 14 May 2024 19:14:56 +0000",
21 | "To": "test@gmail.com",
22 | }
23 | ]
24 |
25 |
26 | # Task storage supporting only a single instance of BabyAGI
27 | class SingleTaskListStorage:
28 | def __init__(self):
29 | self.tasks = deque([])
30 | self.task_id_counter = 0
31 |
32 | def append(self, task: Dict):
33 | self.tasks.append(task)
34 |
35 | def replace(self, tasks: List[Dict]):
36 | self.tasks = deque(tasks)
37 |
38 | def popleft(self):
39 | return self.tasks.popleft()
40 |
41 | def is_empty(self):
42 | return False if self.tasks else True
43 |
44 | def next_task_id(self):
45 | self.task_id_counter += 1
46 | return self.task_id_counter
47 |
48 | def get_task_names(self):
49 | return [t["task_name"] for t in self.tasks]
50 |
51 |
52 | # Initialize tasks storage
53 | tasks_storage = SingleTaskListStorage()
54 |
55 |
56 | def get_ollama_embedding(text):
57 | text = text.replace("\n", " ")
58 | response = ollama.(model="mxbai-embed-large", prompt=text)
59 | return response["embedding"]
60 |
61 |
62 | def task_creation_agent(
63 | objective: str, result: Dict, task_description: str, task_list: List[str]
64 | ):
65 | prompt = f"""
66 | You are to use the result from an execution agent to create new tasks with the following objective: {objective}.
67 | The last completed task has the result: \n{result["data"]}
68 | This result was based on this task description: {task_description}.\n"""
69 |
70 | if task_list:
71 | prompt += f"These are incomplete tasks: {', '.join(task_list)}\n"
72 | prompt += "Based on the result, return a list of tasks to be completed in order to meet the objective. "
73 | if task_list:
74 | prompt += "These new tasks must not overlap with incomplete tasks. "
75 |
76 | prompt += """
77 | Return one task per line in your response. The result must be a numbered list in the format:
78 |
79 | #. First task
80 | #. Second task
81 |
82 | The number of each entry must be followed by a period. If your list is empty, write "There are no tasks to add at this time."
83 | Unless your list is empty, do not include any headers before your numbered list or follow your numbered list with any other output."""
84 |
85 | print(f"\n*****TASK CREATION AGENT PROMPT****\n{prompt}\n")
86 | response = ollama.generate(
87 | model="llama3",
88 | prompt=prompt,
89 | )
90 |
91 | if isinstance(response, dict):
92 | if "response" in response:
93 | response_text = response["response"]
94 | print(f"\n****TASK CREATION AGENT RESPONSE****\n{response_text}\n")
95 | new_tasks = response_text.split("\n")
96 | new_tasks_list = []
97 | for task_string in new_tasks:
98 | task_parts = task_string.strip().split(".", 1)
99 | if len(task_parts) == 2:
100 | task_id = "".join(s for s in task_parts[0] if s.isnumeric())
101 | task_name = re.sub(r"[^\w\s_]+", "", task_parts[1]).strip()
102 | if task_name.strip() and task_id.isnumeric():
103 | new_tasks_list.append(task_name)
104 | out = [{"task_name": task_name} for task_name in new_tasks_list]
105 | return out
106 | else:
107 | raise Exception("No 'response' found in the API response")
108 | else:
109 | raise Exception("Response is not a dictionary")
110 |
111 |
112 | def prioritization_agent():
113 | task_names = tasks_storage.get_task_names()
114 | bullet_string = "\n"
115 |
116 | prompt = f"""
117 | You are tasked with prioritizing the following tasks: {bullet_string + bullet_string.join(task_names)}
118 | Consider the ultimate objective of your team: {OBJECTIVE}.
119 | Tasks should be sorted from highest to lowest priority, where higher-priority tasks are those that act as pre-requisites or are more essential for meeting the objective.
120 | Do not remove any tasks. Return the ranked tasks as a numbered list in the format:
121 |
122 | #. First task
123 | #. Second task
124 |
125 | The entries must be consecutively numbered, starting with 1. The number of each entry must be followed by a period.
126 | Do not include any headers before your ranked list or follow your list with any other output."""
127 |
128 | print(f"\n****TASK PRIORITIZATION AGENT PROMPT****\n{prompt}\n")
129 | response = ollama.generate(
130 | model="llama3",
131 | prompt=prompt,
132 | )
133 |
134 | if isinstance(response, dict):
135 | if "response" in response:
136 | response_text = response["response"]
137 | new_tasks = response_text.strip().split("\n")
138 | else:
139 | raise Exception(f"Unexpected response structure: {response}")
140 | else:
141 | raise Exception("Response is not a dictionary")
142 | print(f"\n****TASK PRIORITIZATION AGENT RESPONSE****\n{response}\n")
143 | if not response:
144 | print(
145 | "Received empty response from prioritization agent. Keeping task list unchanged."
146 | )
147 | return
148 | new_tasks = response_text.split("\n") if "\n" in response_text else [response_text]
149 | new_tasks_list = []
150 | for task_string in new_tasks:
151 | task_parts = task_string.strip().split(".", 1)
152 | if len(task_parts) == 2:
153 | task_id = "".join(s for s in task_parts[0] if s.isnumeric())
154 | task_name = re.sub(r"[^\w\s_]+", "", task_parts[1]).strip()
155 | if task_name.strip():
156 | new_tasks_list.append({"task_id": task_id, "task_name": task_name})
157 |
158 | return new_tasks_list
159 |
160 |
161 | def execution_agent(db, objective: str, task: str) -> str:
162 | context = context_agent(db, query=objective, top_results_num=5)
163 | prompt = f"Perform one task based on the following objective: {objective}.\n"
164 | if context:
165 | prompt += "Take into account these previously completed tasks:" + "\n".join(
166 | context
167 | )
168 | prompt += f"\nYour task: {task}\nResponse:"
169 | response = ollama.generate(
170 | model="llama3",
171 | prompt=f"You are an AI who performs one task based on the following objective: {objective}. Your task: {task}\nResponse:",
172 | stream=False,
173 | )
174 |
175 | if isinstance(response, dict):
176 | if "response" in response:
177 | response_text = response["response"]
178 | new_tasks = response_text.strip().split("\n")
179 | return [{"task_name": task_name} for task_name in new_tasks]
180 | else:
181 | raise Exception(f"Unexpected response structure: {response}")
182 | else:
183 | raise Exception("Response is not a dictionary")
184 |
185 |
186 | def context_agent(db, query: str, top_results_num: int):
187 | query_embedding = get_ollama_embedding(query)
188 | results = db.vector_search(
189 | query_vector=query_embedding, number_of_results=top_results_num
190 | )
191 | results = json.loads(results)
192 | print(f"\n\nContext search results:\n{results}\n\n")
193 | return [row[1].strip('"') for row in results.get("rows", [])]
194 |
195 |
196 | def store_results(db, task: Dict, result: str, result_id: str):
197 | vector = get_ollama_embedding(result)
198 | db.insert_with_vector(
199 | relation_name="tasks",
200 | task_id=result_id,
201 | text=result,
202 | embeddings=vector,
203 | metadata={"task": task["task_name"], "result": result},
204 | )
205 |
206 |
207 | def objective_agent(to, from_email, subject, timestamp, body, attachments):
208 | prompt = f"""
209 | You are an AI assistant that processes emails. You have received an email with the following details:
210 | To: {to}, From: {from_email}, Subject: {subject}, Timestamp: {timestamp}, Body: {body}, Attachments: {attachments}.
211 | Based on this information, determine if the email contains any tasks for the recipient, if any, and return it as a string.
212 | If you don't believe there are any tasks, return the string, "No tasks found." Do not include quotes. RETURN ONLY THIS STRING AND DO NOT INCLUDE ANY OTHER OUTPUT.
213 | """
214 |
215 | print(prompt)
216 | response = ollama.generate(
217 | model="llama3",
218 | prompt=prompt,
219 | )
220 |
221 | # Print the full response for debugging
222 | print(f"Full response: {response}")
223 |
224 | if isinstance(response, dict):
225 | if "response" in response:
226 | response_text = response["response"].strip()
227 | if response_text == "No tasks found." or not response_text:
228 | return {"tasks_found": False, "tasks": []}
229 | else:
230 | new_tasks = response_text.split("\n")
231 | return {
232 | "tasks_found": True,
233 | "tasks": [{"task_name": task_name} for task_name in new_tasks],
234 | }
235 | else:
236 | raise Exception(f"Unexpected response structure: {response}")
237 | else:
238 | raise Exception("Response is not a dictionary")
239 |
240 |
241 | def main():
242 | # Load environment variables from .env file
243 | load_dotenv()
244 |
245 | # Extract email information
246 | email = email_data[0]
247 | to = email["To"]
248 | from_email = email["From"]
249 | subject = email["Subject"]
250 | timestamp = email["Timestamp"]
251 | body = email["Body"]
252 | attachments = "" # Assuming no attachments for simplicity
253 |
254 | # Determine the objective dynamically
255 | objective_response = objective_agent(
256 | to, from_email, subject, timestamp, body, attachments
257 | )
258 |
259 | if not objective_response["tasks_found"]:
260 | print("No tasks identified in the email. Exiting.")
261 | return
262 |
263 | # Set the first task as the objective
264 | OBJECTIVE = objective_response["tasks"][0]["task_name"]
265 | print("\033[96m\033[1m" + "\n*****OBJECTIVE*****\n" + "\033[0m\033[0m")
266 | print(OBJECTIVE)
267 |
268 | # Initialize NexusDB
269 | db = NexusDB()
270 |
271 | JOIN_EXISTING_OBJECTIVE = False
272 |
273 | # Add the initial task if starting new objective
274 | if not JOIN_EXISTING_OBJECTIVE:
275 | initial_task = {
276 | "task_id": tasks_storage.next_task_id(),
277 | "task_name": "Develop a task list.",
278 | }
279 | tasks_storage.append(initial_task)
280 |
281 | # Main loop
282 | loop = True
283 | while loop:
284 | if not tasks_storage.is_empty():
285 | print("\033[95m\033[1m" + "\n*****TASK LIST*****\n" + "\033[0m\033[0m")
286 | for t in tasks_storage.get_task_names():
287 | print(" • " + str(t))
288 |
289 | task = tasks_storage.popleft()
290 | print("\033[92m\033[1m" + "\n*****NEXT TASK*****\n" + "\033[0m\033[0m")
291 | print(str(task["task_name"]))
292 |
293 | result = execution_agent(db, OBJECTIVE, str(task["task_name"]))
294 | print("\033[93m\033[1m" + "\n*****TASK RESULT*****\n" + "\033[0m\033[0m")
295 | print(result)
296 |
297 | enriched_result = {"data": result}
298 | result_id = f"result_{task['task_id']}"
299 |
300 | store_results(db, task, result, result_id)
301 |
302 | new_tasks = task_creation_agent(
303 | OBJECTIVE,
304 | enriched_result,
305 | task["task_name"],
306 | tasks_storage.get_task_names(),
307 | )
308 |
309 | print("Adding new tasks to task_storage")
310 | for new_task in new_tasks:
311 | new_task.update({"task_id": tasks_storage.next_task_id()})
312 | print(str(new_task))
313 | tasks_storage.append(new_task)
314 |
315 | if not JOIN_EXISTING_OBJECTIVE:
316 | prioritized_tasks = prioritization_agent()
317 | if prioritized_tasks:
318 | tasks_storage.replace(prioritized_tasks)
319 |
320 | time.sleep(5)
321 | else:
322 | print("Done.")
323 | loop = False
324 |
325 |
326 | if __name__ == "__main__":
327 | main()
328 |
--------------------------------------------------------------------------------
/tests/colorlogs.py:
--------------------------------------------------------------------------------
1 | import colorlog
2 | from colorlog import ColoredFormatter
3 |
4 | # Define a mapping from thread names to colors
5 | THREAD_COLOR_MAPPING = {
6 | "EmailProcessor": "red",
7 | "email_fetcher": "green",
8 | "EntityExtraction": "blue",
9 | # Add other thread name to color mappings here
10 | }
11 |
12 |
13 | class ThreadNameColoredFormatter(ColoredFormatter):
14 | def __init__(self, *args, **kwargs):
15 | super().__init__(*args, **kwargs)
16 | # Set the default log colors based on log levels
17 | self.default_log_colors = {
18 | "DEBUG": "cyan",
19 | "INFO": "white", # Default color for INFO, will be overridden by thread color if available
20 | "WARNING": "yellow",
21 | "ERROR": "red",
22 | "CRITICAL": "red,bg_white",
23 | }
24 |
25 | def format(self, record):
26 | # Set the log color based on log level
27 | self.log_colors = self.default_log_colors.copy()
28 |
29 | # Override the log color for INFO based on thread name if applicable
30 | if record.levelname == "INFO":
31 | thread_name = (
32 | record.threadName.split("-")[0] if record.threadName else "Thread"
33 | ) # Use the base name for mapping
34 | thread_log_color = THREAD_COLOR_MAPPING.get(thread_name, "white")
35 | self.log_colors["INFO"] = thread_log_color
36 |
37 | return super().format(record)
38 |
39 |
40 | # Configure the formatter
41 | formatter = ThreadNameColoredFormatter(
42 | "%(log_color)s%(asctime)s - %(name)s - %(levelname)s - [%(threadName)s] - %(message)s",
43 | datefmt="%Y-%m-%d %H:%M:%S",
44 | )
45 |
46 | # Usage example in the main application setup
47 | if __name__ == "__main__":
48 | import logging
49 | import threading
50 |
51 | # Set up the handler and logger
52 | handler = logging.StreamHandler()
53 | handler.setFormatter(formatter)
54 |
55 | logger = logging.getLogger()
56 | logger.addHandler(handler)
57 | logger.setLevel(logging.DEBUG)
58 |
59 | # Define a test function to generate logs in different threads
60 | def log_messages(logger):
61 | logger.debug("This is a DEBUG message")
62 | logger.info("This is an INFO message")
63 | logger.warning("This is a WARNING message")
64 | logger.error("This is an ERROR message")
65 | logger.critical("This is a CRITICAL message")
66 |
67 | # Create and start threads
68 | threads = []
69 | for thread_name in ["EmailProcessor", "email_fetcher", "EntityExtraction"]:
70 | thread_logger = logging.getLogger(thread_name)
71 | thread = threading.Thread(
72 | name=thread_name, target=log_messages, args=(thread_logger,)
73 | )
74 | threads.append(thread)
75 | thread.start()
76 |
77 | # Log from the main thread
78 | log_messages(logger)
79 |
80 | # Wait for all threads to complete
81 | for thread in threads:
82 | thread.join()
83 |
--------------------------------------------------------------------------------
/tests/embedding.py:
--------------------------------------------------------------------------------
1 | import ollama
2 |
3 |
4 | def get_embedding_length():
5 | model = "mxbai-embed-large"
6 | prompt = "NexusDB is the best database."
7 |
8 | response = ollama.(model, prompt)
9 | print(f"Embedding length: {len(response['embedding'])}")
10 |
11 |
12 | def chat():
13 | print("Chatting with the Llama AI.../n/n")
14 | model = "llama3"
15 | prompt = "Please say 'Hello' and nothing else."
16 |
17 | response = ollama.generate(model, prompt, stream=False)
18 | print(response)
19 |
20 |
21 | # Call the function
22 | chat()
23 | # get_embedding_length()
24 |
--------------------------------------------------------------------------------
/tests/entity_add.py:
--------------------------------------------------------------------------------
1 | import json
2 | from typing import List
3 |
4 | import ollama
5 | from ollama import Message
6 |
7 |
8 | def conditional_entity_addition(data):
9 | # Retrieve the entity type from the data
10 | entity_type = data.get("entity_type", None)
11 | # If no entity type is provided, return an error
12 | if not entity_type:
13 | return {"error": "Entity type is required."}, 400
14 |
15 | # Adjusted to access nested 'data'
16 | entity_data = data.get("data", {})
17 |
18 | search_results = [
19 | {"id": "1", "name": "John Smith"},
20 | {"id": "2", "name": "Jane Doe"},
21 | ]
22 |
23 | # Combine all search results
24 | combined_results = {result["id"]: result for result in search_results}.values()
25 | print(f"Combined results: {list(combined_results)}")
26 | combined_results_str = ", ".join(json.dumps(result) for result in combined_results)
27 |
28 | # Prepare the message for OpenAI API
29 | prompt: List[Message] = [
30 | Message(
31 | role="system",
32 | content="You are a helpful assistant who's specialty is to decide if new input data matches data already in our database. Review the search results provided, compare against the input data, and if there's a match respond with the ID number of the match, and only the ID number. If there are no matches, respond with 'No Matches'. Your response is ALWAYS an ID number alone, or 'No Matches'. When reviewing whether a match existings in our search results to our new input, take into account that the name may not match perfectly (for example, one might have just a first name, or a nick name, while the other has a full name), in which case look at the additional information about the user to determine if there's a strong likelihood they are the same person. For companies, you should consider different names of the same company as the same, such as EA and Electronic Arts (make your best guess). If the likelihood is strong, respond with and only with the ID number. If likelihood is low, respond with 'No Matches'.",
33 | ),
34 | Message(
35 | role="user",
36 | content=f"Here are the search results: {combined_results_str}. Does any entry match the input data: {data}?",
37 | ),
38 | ]
39 |
40 | # Make a call to OpenAI API
41 | try:
42 | response = ollama.chat(
43 | model="llama3",
44 | messages=prompt,
45 | stream=True,
46 | )
47 |
48 | if isinstance(response, dict) and "response" in response:
49 | response_text = response["response"]
50 | return response_text
51 | else:
52 | ai_response = ""
53 | for chunk in response:
54 | if isinstance(chunk, dict):
55 | if (
56 | "message" in chunk
57 | and isinstance(chunk["message"], dict)
58 | and "content" in chunk["message"]
59 | ):
60 | print(chunk["message"]["content"], end="", flush=True)
61 | ai_response += chunk["message"]["content"]
62 | else:
63 | raise Exception("Invalid chunk structure")
64 | print(f"AI response: {ai_response}")
65 |
66 | # Process the AI's response
67 | if "no matches" in ai_response.lower():
68 | # If no match found, add the new entity
69 | # entity_id = add_entity(entity_type, data)
70 | print("adding entity\n\n")
71 | entity_id = "123"
72 | return {"success": True, "entity_id": entity_id}, 200
73 | else:
74 | # If a match is found, return the match details
75 | match_id = ai_response
76 | return {
77 | "success": False,
78 | "message": "Match found",
79 | "match_id": match_id,
80 | }, 200
81 |
82 | except Exception as e:
83 | print(f"Error calling OpenAI: {e}")
84 | return {"error": str(e)}, 500
85 |
86 |
87 | # Mocking the input data
88 | input_data = {"entity_type": "person", "data": {"name": "John Doe", "age": 30}}
89 |
90 | # Call the function and print the result
91 | result, status_code = conditional_entity_addition(input_data)
92 | print(result, status_code)
93 |
--------------------------------------------------------------------------------
/tests/graph_agent.py:
--------------------------------------------------------------------------------
1 | import ollama
2 |
3 |
4 | def entity_extraction_agent(text_input):
5 | prompt = [
6 | {
7 | "role": "system",
8 | "content": """You are an AI expert specializing in knowledge graph creation with the goal of capturing relationships based on a given input or request.
9 | You are given input in various forms such as paragraph, email, text files, and more.
10 | Your task is to create a knowledge graph based on the input.
11 | Only use organizations, people, and events as nodes and do not include concepts or products.
12 | Only add nodes that have a relationship with at least one other node.
13 | Make sure that the node type (people, org, event) matches the to_type or for_type when the entity is part of a relationship.
14 | Return the knowledge graph as a JSON object. DO NOT INCLUDE ANYTHING ELSE IN THE RESPONSE.""",
15 | },
16 | {
17 | "role": "user",
18 | "content": "Can you please help John Smith from IT get access to the system? He needs it as part of the IT Modernization effort.",
19 | },
20 | {
21 | "role": "assistant",
22 | "content": '{"entities": [{"name": "Modernization of the IT infrastructure", "type": "Project", "description": "A project to modernize the IT infrastructure of the company.", "department": "IT",},{"name": "Person A", "type": "Person", "memberOf": "IT",},{"name": "IT", "type": "Organization", "description": "The IT department of the company.", "member": "Person A",},]}',
23 | },
24 | {"role": "user", "content": text_input},
25 | ]
26 |
27 | response = ollama.chat(
28 | model="llama3",
29 | messages=prompt,
30 | stream=True,
31 | )
32 |
33 | if isinstance(response, dict) and "response" in response:
34 | response_text = response["response"]
35 | return response_text
36 | else:
37 | try:
38 | for chunk in response:
39 | if isinstance(chunk, dict):
40 | if (
41 | "message" in chunk
42 | and isinstance(chunk["message"], dict)
43 | and "content" in chunk["message"]
44 | ):
45 | print(chunk["message"]["content"], end="", flush=True)
46 | else:
47 | raise Exception("Invalid chunk structure")
48 | return "done"
49 | except Exception as e:
50 | raise Exception(f"No 'response' found in the API response: {e}")
51 |
52 |
53 | response = entity_extraction_agent(
54 | "Adam from team A will be able to help answer any questions."
55 | )
56 | print(response)
57 |
--------------------------------------------------------------------------------
/tests/ollama_raw.py:
--------------------------------------------------------------------------------
1 | import ollama
2 |
3 |
4 | def entity_extraction_agent(text_input):
5 | response = ollama.generate(
6 | model="llama3",
7 | prompt=text_input,
8 | # raw=True,
9 | stream=True,
10 | )
11 |
12 | if isinstance(response, dict) and "response" in response:
13 | response_text = response["response"]
14 | return response_text
15 | else:
16 | try:
17 | for chunk in response:
18 | if isinstance(chunk, dict):
19 | print(chunk["response"], end="", flush=True)
20 | return "done"
21 | except Exception as e:
22 | raise Exception(f"No 'response' found in the API response: {e}")
23 |
24 |
25 | response = entity_extraction_agent(".")
26 | print(response)
27 |
--------------------------------------------------------------------------------
/tests/ollama_streaming.py:
--------------------------------------------------------------------------------
1 | import threading
2 | from typing import List
3 |
4 | from ollama import Message
5 |
6 | from utils.ollama import ollama_chat
7 |
8 | # Dummy Messages for Testing
9 | messages: List[Message] = [
10 | Message(role="user", content="Hello, how are you?"),
11 | Message(role="user", content="What is the weather like today?"),
12 | Message(role="user", content="Tell me a joke."),
13 | Message(role="user", content="What's the capital of France?"),
14 | Message(role="user", content="What's the latest news?"),
15 | ]
16 |
17 | # Dummy model name
18 | model = "llama3"
19 |
20 | # Define the number of threads
21 | num_threads = len(messages)
22 |
23 |
24 | # Test function to run in threads
25 | def test_ollama_chat(thread_id: int, message: Message):
26 | print(f"Thread-{thread_id} starting with message: {message}")
27 | response = ollama_chat(model=model, messages=[message], stream=True)
28 | print(f"\nThread-{thread_id} received response:\n{response}\n")
29 |
30 |
31 | # Create and start threads
32 | threads = []
33 | for i in range(num_threads):
34 | thread = threading.Thread(target=test_ollama_chat, args=(i, messages[i]))
35 | threads.append(thread)
36 | thread.start()
37 |
38 | # Wait for all threads to complete
39 | for thread in threads:
40 | thread.join()
41 |
42 | print("All threads completed.")
43 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Astra-Analytics/Task-Agent-Starter-Kit/011cccfd852920748f66d7e6dd80ca2aaf974f70/utils/__init__.py
--------------------------------------------------------------------------------
/utils/custom_log_formatter.py:
--------------------------------------------------------------------------------
1 | from colorlog import ColoredFormatter
2 |
3 | # Define a mapping from thread names to colors
4 | THREAD_COLOR_MAPPING = {
5 | "EmailProcessor": "cyan",
6 | "email_fetcher": "purple",
7 | "EntityExtraction": "blue",
8 | # Add other thread name to color mappings here
9 | }
10 |
11 |
12 | class ThreadNameColoredFormatter(ColoredFormatter):
13 | def __init__(self, *args, **kwargs):
14 | super().__init__(*args, **kwargs)
15 | # Set the default log colors based on log levels
16 | self.default_log_colors = {
17 | "DEBUG": "light_white",
18 | "INFO": "white", # Default color for INFO, will be overridden by thread color if available
19 | "WARNING": "yellow",
20 | "ERROR": "red",
21 | "CRITICAL": "red,bg_white",
22 | }
23 |
24 | def format(self, record):
25 | # Set the log color based on log level
26 | self.log_colors = self.default_log_colors.copy()
27 |
28 | # Override the log color for INFO based on thread name if applicable
29 | if record.levelname == "INFO":
30 | thread_name = (
31 | record.threadName.split("-")[0] if record.threadName else "Thread"
32 | ) # Use the base name for mapping
33 | thread_log_color = THREAD_COLOR_MAPPING.get(thread_name, "white")
34 | self.log_colors["INFO"] = thread_log_color
35 |
36 | return super().format(record)
37 |
--------------------------------------------------------------------------------
/utils/ollama.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import threading
3 | from typing import Any, Dict, Iterator, List, Mapping, Union
4 |
5 | import ollama
6 | from ollama import Message
7 |
8 | logger = logging.getLogger(__name__)
9 |
10 | # Initialize a lock
11 | print_lock = threading.Lock()
12 |
13 |
14 | def get_ollama_embedding(text):
15 | text = text.replace("\n", " ")
16 | response = ollama.embeddings(model="mxbai-embed-large", prompt=text)
17 | return response["embedding"]
18 |
19 |
20 | def handle_response(
21 | response: Union[Dict[str, Any], Iterator[Mapping[str, Any]]], stream: bool = False
22 | ) -> str:
23 | if isinstance(response, dict) and "response" in response:
24 | return response["response"].strip()
25 | elif stream:
26 | ai_response = ""
27 | try:
28 | with print_lock: # Acquire the lock
29 | for chunk in response:
30 | if isinstance(chunk, Mapping) and "message" in chunk:
31 | message = chunk["message"]
32 | if isinstance(message, Mapping) and "content" in message:
33 | print(message["content"], end="", flush=True)
34 | ai_response += message["content"]
35 | elif isinstance(message, str):
36 | print(message, end="", flush=True)
37 | ai_response += message
38 | else:
39 | raise Exception("Invalid chunk structure")
40 | elif isinstance(chunk, Mapping) and "response" in chunk:
41 | print(chunk["response"], end="", flush=True)
42 | ai_response += chunk["response"]
43 | else:
44 | raise Exception("Invalid chunk structure")
45 | return ai_response
46 | except Exception as e:
47 | raise Exception(f"No 'response' found in the API response: {e}")
48 | else:
49 | raise Exception(f"Unexpected response structure: {response}")
50 |
51 |
52 | def ollama_generate(model: str, prompt: str, stream: bool = False) -> str:
53 | response = ollama.generate(model=model, prompt=prompt, stream=stream)
54 | if isinstance(response, (dict, Iterator)):
55 | return handle_response(response, stream=stream)
56 | else:
57 | raise TypeError("Invalid response type")
58 |
59 |
60 | def ollama_chat(model: str, messages: List[Message], stream: bool = False) -> str:
61 | response = ollama.chat(model=model, messages=messages, stream=stream)
62 | if isinstance(response, (dict, Iterator)):
63 | return handle_response(response, stream=stream)
64 | else:
65 | raise TypeError("Invalid response type")
66 |
--------------------------------------------------------------------------------