├── .env.ci.example ├── .github └── workflows │ ├── CI.yml │ └── publish_to_pypi.yml ├── .gitignore ├── .pre-commit-config.yaml ├── LICENSE ├── README.md ├── examples ├── .env.example ├── attachment.py ├── dataset.py ├── file-upload.py ├── langchain_chatopenai.py ├── langchain_toolcall.py ├── langchain_variable.py ├── llamaindex.py ├── llamaindex_workflow.py ├── main.py ├── multimodal.py ├── openai_agents.py ├── parallel_requests.py ├── prompt.py ├── samplefile.txt ├── streaming.py ├── thread.py └── user.py ├── literalai ├── __init__.py ├── api │ ├── README.md │ ├── __init__.py │ ├── asynchronous.py │ ├── base.py │ ├── helpers │ │ ├── __init__.py │ │ ├── attachment_helpers.py │ │ ├── dataset_helpers.py │ │ ├── generation_helpers.py │ │ ├── gql.py │ │ ├── prompt_helpers.py │ │ ├── score_helpers.py │ │ ├── step_helpers.py │ │ ├── thread_helpers.py │ │ └── user_helpers.py │ └── synchronous.py ├── cache │ ├── __init__.py │ ├── prompt_helpers.py │ └── shared_cache.py ├── callback │ ├── __init__.py │ ├── langchain_callback.py │ └── openai_agents_processor.py ├── client.py ├── context.py ├── environment.py ├── evaluation │ ├── __init__.py │ ├── dataset.py │ ├── dataset_experiment.py │ ├── dataset_item.py │ └── experiment_item_run.py ├── event_processor.py ├── exporter.py ├── helper.py ├── instrumentation │ ├── __init__.py │ ├── llamaindex │ │ ├── __init__.py │ │ ├── event_handler.py │ │ └── span_handler.py │ ├── mistralai.py │ └── openai.py ├── my_types.py ├── observability │ ├── __init__.py │ ├── filter.py │ ├── generation.py │ ├── message.py │ ├── step.py │ └── thread.py ├── prompt_engineering │ ├── __init__.py │ └── prompt.py ├── py.typed ├── requirements.py ├── version.py └── wrappers.py ├── mypy.ini ├── pytest.ini ├── requirements-dev.txt ├── requirements.txt ├── run-test.sh ├── setup.py └── tests ├── __init__.py ├── e2e ├── test_e2e.py ├── test_llamaindex.py ├── test_mistralai.py └── test_openai.py └── unit ├── __init__.py └── test_cache.py /.env.ci.example: -------------------------------------------------------------------------------- 1 | LITERAL_API_URL= 2 | LITERAL_API_KEY= 3 | -------------------------------------------------------------------------------- /.github/workflows/CI.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | workflow_dispatch: 5 | pull_request: 6 | branches: ["main"] 7 | push: 8 | branches: 9 | - main 10 | 11 | permissions: 12 | contents: read 13 | 14 | jobs: 15 | lint: 16 | runs-on: ubuntu-latest 17 | steps: 18 | - uses: actions/checkout@v3 19 | - name: Set up Python 3.9 20 | uses: actions/setup-python@v3 21 | with: 22 | python-version: "3.9" 23 | - name: Install dependencies 24 | run: | 25 | python -m pip install --upgrade pip 26 | if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi 27 | if [ -f requirements.txt ]; then pip install -r requirements.txt; fi 28 | - name: Lint and format with ruff 29 | run: | 30 | ruff check 31 | - name: Type check 32 | run: | 33 | mypy . 34 | e2e-tests: 35 | runs-on: ubuntu-latest 36 | steps: 37 | - uses: actions/checkout@v3 38 | - name: Set up Python 3.9 39 | uses: actions/setup-python@v3 40 | with: 41 | python-version: "3.9" 42 | - name: Install dependencies 43 | run: | 44 | python -m pip install --upgrade pip 45 | pip install . 46 | if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi 47 | if [ -f requirements.txt ]; then pip install -r requirements.txt; fi 48 | - name: Test with pytest 49 | env: 50 | LITERAL_API_URL: ${{ secrets.LITERAL_API_URL }} 51 | LITERAL_API_KEY: ${{ secrets.LITERAL_API_KEY }} 52 | run: | 53 | pytest -m e2e 54 | -------------------------------------------------------------------------------- /.github/workflows/publish_to_pypi.yml: -------------------------------------------------------------------------------- 1 | name: Publish 2 | 3 | on: 4 | workflow_dispatch: 5 | release: 6 | types: [published] 7 | 8 | jobs: 9 | build-n-publish: 10 | name: Upload release to PyPI 11 | runs-on: ubuntu-latest 12 | env: 13 | name: pypi 14 | url: https://pypi.org/p/literalai 15 | permissions: 16 | contents: read 17 | id-token: write 18 | steps: 19 | - uses: actions/checkout@v3 20 | with: 21 | ref: main 22 | - name: Set up Python 23 | uses: actions/setup-python@v4 24 | with: 25 | python-version: "3.9" 26 | - name: Install dependencies and build 27 | run: | 28 | python -m pip install --upgrade pip 29 | pip install -r requirements.txt 30 | python setup.py sdist 31 | - name: Publish package distributions to PyPI 32 | uses: pypa/gh-action-pypi-publish@release/v1 33 | with: 34 | packages-dir: dist 35 | password: ${{ secrets.PYPI_API_TOKEN }} 36 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | 3 | .DS_Store 4 | 5 | # Byte-compiled / optimized / DLL files 6 | __pycache__/ 7 | *.py[cod] 8 | *$py.class 9 | 10 | # C extensions 11 | *.so 12 | .vscode/ 13 | 14 | # Distribution / packaging 15 | .Python 16 | build/ 17 | develop-eggs/ 18 | dist/ 19 | downloads/ 20 | eggs/ 21 | .eggs/ 22 | lib/ 23 | lib64/ 24 | parts/ 25 | sdist/ 26 | var/ 27 | wheels/ 28 | share/python-wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .nox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | .ruff_cache 52 | nosetests.xml 53 | coverage.xml 54 | *.cover 55 | *.py,cover 56 | .hypothesis/ 57 | .pytest_cache/ 58 | cover/ 59 | 60 | # Translations 61 | *.mo 62 | *.pot 63 | 64 | # Django stuff: 65 | *.log 66 | local_settings.py 67 | db.sqlite3 68 | db.sqlite3-journal 69 | 70 | # Flask stuff: 71 | instance/ 72 | .webassets-cache 73 | 74 | # Scrapy stuff: 75 | .scrapy 76 | 77 | # Sphinx documentation 78 | docs/_build/ 79 | 80 | # PyBuilder 81 | .pybuilder/ 82 | target/ 83 | 84 | # Jupyter Notebook 85 | .ipynb_checkpoints 86 | 87 | # IPython 88 | profile_default/ 89 | ipython_config.py 90 | 91 | # pyenv 92 | # For a library or package, you might want to ignore these files since the code is 93 | # intended to run in multiple environments; otherwise, check them in: 94 | .python-version 95 | 96 | # pipenv 97 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 98 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 99 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 100 | # install all needed dependencies. 101 | #Pipfile.lock 102 | 103 | # poetry 104 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 105 | # This is especially recommended for binary packages to ensure reproducibility, and is more 106 | # commonly ignored for libraries. 107 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 108 | #poetry.lock 109 | 110 | # pdm 111 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 112 | #pdm.lock 113 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 114 | # in version control. 115 | # https://pdm.fming.dev/#use-with-ide 116 | .pdm.toml 117 | 118 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 119 | __pypackages__/ 120 | 121 | # Celery stuff 122 | celerybeat-schedule 123 | celerybeat.pid 124 | 125 | # SageMath parsed files 126 | *.sage.py 127 | 128 | # Environments 129 | .env 130 | .venv 131 | env/ 132 | venv/ 133 | py_examples/ 134 | ENV/ 135 | env.bak/ 136 | venv.bak/ 137 | 138 | # Spyder project settings 139 | .spyderproject 140 | .spyproject 141 | 142 | # Rope project settings 143 | .ropeproject 144 | 145 | # mkdocs documentation 146 | /site 147 | 148 | # mypy 149 | .mypy_cache/ 150 | .dmypy.json 151 | dmypy.json 152 | 153 | # Pyre type checker 154 | .pyre/ 155 | 156 | # pytype static type analyzer 157 | .pytype/ 158 | 159 | # Cython debug symbols 160 | cython_debug/ 161 | 162 | # PyCharm 163 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 164 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 165 | # and can be added to the global gitignore or merged into this file. For a more nuclear 166 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 167 | .idea/ 168 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/psf/black 3 | rev: 23.11.0 4 | hooks: 5 | - id: black 6 | - repo: https://github.com/pre-commit/mirrors-isort 7 | rev: "v5.10.1" # Use the revision sha / tag you want to point at 8 | hooks: 9 | - id: isort 10 | args: ["--profile", "black", "--filter-files"] 11 | - repo: https://github.com/astral-sh/ruff-pre-commit 12 | rev: v0.1.11 13 | hooks: 14 | - id: ruff 15 | - repo: https://github.com/pre-commit/mirrors-mypy 16 | rev: "v1.7.1" # Use the sha / tag you want to point at 17 | hooks: 18 | - id: mypy 19 | additional_dependencies: 20 | - pydantic 21 | - httpx 22 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Literal AI client 2 | 3 | ## Installation 4 | 5 | ```bash 6 | pip install literalai 7 | ``` 8 | 9 | ## Usage 10 | 11 | The full documentation is available [here](https://docs.getliteral.ai/python-client). 12 | 13 | Create a `.env` file with the `LITERAL_API_KEY` environment variable set to your API key. 14 | 15 | ```python 16 | from literalai import LiteralClient 17 | from dotenv import load_dotenv 18 | 19 | load_dotenv() 20 | 21 | literalai_client = LiteralClient() 22 | 23 | @literalai_client.step(type="run") 24 | def my_step(input): 25 | return f"World" 26 | 27 | 28 | @literalai_client.thread 29 | def main(): 30 | print(my_step("Hello")) 31 | 32 | 33 | main() 34 | client.flush_and_stop() 35 | print("Done") 36 | ``` 37 | 38 | ## Development setup 39 | 40 | ```bash 41 | pip install -r requirements-dev.txt 42 | ``` 43 | -------------------------------------------------------------------------------- /examples/.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY="sk-" 2 | TAVILY_API_KEY="" 3 | LITERAL_API_KEY="lsk_" 4 | LITERAL_API_URL="http://localhost:3000" -------------------------------------------------------------------------------- /examples/attachment.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | 5 | from literalai import LiteralClient 6 | 7 | load_dotenv() 8 | 9 | 10 | sdk = LiteralClient(batch_size=2) 11 | 12 | 13 | async def main(): 14 | thread = sdk.api.create_thread() 15 | step = sdk.start_step(name="test", thread_id=thread.id) 16 | sdk.api.send_steps([step]) 17 | 18 | try: 19 | attachment = sdk.api.create_attachment( 20 | name="test", 21 | url="https://www.perdu.com/", 22 | mime="text/html", 23 | thread_id=thread.id, 24 | step_id=step.id, 25 | ) 26 | 27 | print(attachment.to_dict()) 28 | 29 | attachment = sdk.api.update_attachment( 30 | id=attachment.id, 31 | update_params={ 32 | "url": "https://api.github.com/repos/chainlit/chainlit", 33 | "mime": "application/json", 34 | "metadata": {"test": "test"}, 35 | }, 36 | ) 37 | 38 | print(attachment.to_dict()) 39 | 40 | attachment = sdk.api.get_attachment(id=attachment.id) 41 | 42 | print(attachment.to_dict()) 43 | 44 | sdk.api.delete_attachment(id=attachment.id) 45 | 46 | try: 47 | attachment = sdk.api.get_attachment(id=attachment.id) 48 | except Exception as e: 49 | print(e) 50 | 51 | finally: 52 | sdk.api.delete_thread(id=thread.id) 53 | 54 | 55 | asyncio.run(main()) 56 | -------------------------------------------------------------------------------- /examples/dataset.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | 5 | from literalai import AsyncLiteralClient, LiteralClient 6 | 7 | load_dotenv() 8 | 9 | sdk = LiteralClient(batch_size=2) 10 | async_sdk = AsyncLiteralClient(batch_size=2) 11 | 12 | 13 | async def init(): 14 | thread = await async_sdk.api.create_thread() 15 | return await async_sdk.api.create_step( 16 | thread_id=thread.id, 17 | input={"content": "hello world!"}, 18 | output={"content": "hello back!"}, 19 | ) 20 | 21 | 22 | step = asyncio.run(init()) 23 | 24 | 25 | async def main_async(): 26 | # Create a dataset 27 | dataset = await async_sdk.api.create_dataset( 28 | name="foo", description="bar", metadata={"demo": True} 29 | ) 30 | assert dataset.name == "foo" 31 | assert dataset.description == "bar" 32 | assert dataset.metadata == {"demo": True} 33 | 34 | # Update a dataset 35 | dataset = await async_sdk.api.update_dataset(id=dataset.id, name="baz") 36 | assert dataset.name == "baz" 37 | 38 | # Create dataset items 39 | items = [ 40 | { 41 | "input": {"content": "What is literal?"}, 42 | "expected_output": {"content": "Literal is an observability solution."}, 43 | }, 44 | { 45 | "input": {"content": "How can I install the sdk?"}, 46 | "expected_output": {"content": "pip install literalai"}, 47 | }, 48 | ] 49 | items = [ 50 | await async_sdk.api.create_dataset_item(dataset_id=dataset.id, **item) 51 | for item in items 52 | ] 53 | 54 | for item in items: 55 | assert item.dataset_id == dataset.id 56 | assert item.input is not None 57 | assert item.expected_output is not None 58 | 59 | # Get a dataset with items 60 | dataset = await async_sdk.api.get_dataset(id=dataset.id) 61 | 62 | assert dataset.items is not None 63 | assert len(dataset.items) == 2 64 | 65 | # Add step to dataset 66 | assert step.id is not None 67 | step_item = await async_sdk.api.add_step_to_dataset(dataset.id, step.id) 68 | 69 | assert step_item.input == {"content": "hello world!"} 70 | assert step_item.expected_output == {"content": "hello back!"} 71 | 72 | # Delete a dataset item 73 | await async_sdk.api.delete_dataset_item(id=items[0].id) 74 | 75 | # Delete a dataset 76 | await async_sdk.api.delete_dataset(id=dataset.id) 77 | 78 | 79 | def main_sync(): 80 | # Create a dataset 81 | dataset = sdk.api.create_dataset( 82 | name="foo", description="bar", metadata={"demo": True} 83 | ) 84 | assert dataset.name == "foo" 85 | assert dataset.description == "bar" 86 | assert dataset.metadata == {"demo": True} 87 | 88 | # Update a dataset 89 | dataset = sdk.api.update_dataset(id=dataset.id, name="baz") 90 | assert dataset.name == "baz" 91 | 92 | # Create dataset items 93 | items = [ 94 | { 95 | "input": {"content": "What is literal?"}, 96 | "expected_output": {"content": "Literal is an observability solution."}, 97 | }, 98 | { 99 | "input": {"content": "How can I install the sdk?"}, 100 | "expected_output": {"content": "pip install literalai"}, 101 | }, 102 | ] 103 | items = [ 104 | sdk.api.create_dataset_item(dataset_id=dataset.id, **item) for item in items 105 | ] 106 | 107 | for item in items: 108 | assert item.dataset_id == dataset.id 109 | assert item.input is not None 110 | assert item.expected_output is not None 111 | 112 | # Get a dataset with items 113 | dataset = sdk.api.get_dataset(id=dataset.id) 114 | 115 | assert dataset.items is not None 116 | assert len(dataset.items) == 2 117 | 118 | # Add step to dataset 119 | assert step.id is not None 120 | step_item = sdk.api.add_step_to_dataset(dataset.id, step.id) 121 | 122 | assert step_item.input == {"content": "hello world!"} 123 | assert step_item.expected_output == {"content": "hello back!"} 124 | 125 | # Delete a dataset item 126 | sdk.api.delete_dataset_item(id=items[0].id) 127 | 128 | # Delete a dataset 129 | sdk.api.delete_dataset(id=dataset.id) 130 | 131 | 132 | asyncio.run(main_async()) 133 | 134 | main_sync() 135 | -------------------------------------------------------------------------------- /examples/file-upload.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import mimetypes 3 | from pathlib import Path 4 | 5 | from dotenv import load_dotenv 6 | 7 | from literalai import LiteralClient 8 | 9 | load_dotenv() 10 | 11 | 12 | sdk = LiteralClient(batch_size=2) 13 | 14 | 15 | async def main(): 16 | thread = sdk.api.create_thread(metadata={"key": "value"}, tags=["hello"]) 17 | 18 | id = thread.id 19 | path = Path(__file__).parent / "./samplefile.txt" 20 | mime, _ = mimetypes.guess_type(path) 21 | 22 | with open(path, "rb") as file: 23 | data = file.read() 24 | 25 | res = sdk.api.upload_file(content=data, mime=mime, thread_id=id) 26 | 27 | print(res) 28 | 29 | 30 | asyncio.run(main()) 31 | -------------------------------------------------------------------------------- /examples/langchain_chatopenai.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | from langchain.schema import HumanMessage 5 | from langchain_community.chat_models import ChatOpenAI # type: ignore 6 | 7 | from literalai import LiteralClient 8 | 9 | load_dotenv() 10 | 11 | client = LiteralClient() 12 | chat_model = ChatOpenAI() 13 | 14 | 15 | @client.thread(name="main") 16 | async def main(): 17 | text = "What would be a good company name for a company that makes colorful socks?" 18 | messages = [HumanMessage(content=text)] 19 | 20 | cb = client.langchain_callback() 21 | with client.step(name="chat_model.invoke"): 22 | print(chat_model.invoke(messages, config={"callbacks": [cb], "tags": ["test"]})) 23 | 24 | print( 25 | ( 26 | await chat_model.ainvoke( 27 | messages, config={"callbacks": [client.langchain_callback()]} 28 | ) 29 | ) 30 | ) 31 | 32 | 33 | asyncio.run(main()) 34 | client.flush_and_stop() 35 | 36 | print("Done") 37 | -------------------------------------------------------------------------------- /examples/langchain_toolcall.py: -------------------------------------------------------------------------------- 1 | from dotenv import load_dotenv 2 | from langchain.agents import AgentExecutor, create_tool_calling_agent 3 | from langchain.agents.agent import BaseSingleActionAgent 4 | from langchain_community.tools.tavily_search import TavilySearchResults 5 | from langchain_core.messages import AIMessage, HumanMessage 6 | from langchain_core.runnables.config import RunnableConfig 7 | from langchain_openai import ChatOpenAI # type: ignore 8 | 9 | from literalai import LiteralClient 10 | 11 | # Add OPENAI_API_KEY and TAVILY_API_KEY for this example. 12 | load_dotenv() 13 | 14 | model = ChatOpenAI(model="gpt-4o") 15 | search = TavilySearchResults(max_results=2) 16 | tools = [search] 17 | 18 | lai_client = LiteralClient() 19 | lai_client.initialize() 20 | lai_prompt = lai_client.api.get_or_create_prompt( 21 | name="LC Agent", 22 | settings={ 23 | "model": "gpt-4o-mini", 24 | "top_p": 1, 25 | "provider": "openai", 26 | "max_tokens": 4095, 27 | "temperature": 0, 28 | "presence_penalty": 0, 29 | "frequency_penalty": 0, 30 | }, 31 | template_messages=[ 32 | {"role": "system", "content": "You are a helpful assistant"}, 33 | {"role": "assistant", "content": "{{chat_history}}"}, 34 | {"role": "user", "content": "{{input}}"}, 35 | {"role": "assistant", "content": "{{agent_scratchpad}}"}, 36 | ], 37 | ) 38 | prompt = lai_prompt.to_langchain_chat_prompt_template( 39 | additional_messages=[("placeholder", "{agent_scratchpad}")], 40 | ) 41 | 42 | agent: BaseSingleActionAgent = create_tool_calling_agent(model, tools, prompt) # type: ignore 43 | agent_executor = AgentExecutor(agent=agent, tools=tools) 44 | 45 | # Replace with ainvoke for asynchronous execution. 46 | agent_executor.invoke( 47 | { 48 | "chat_history": [ 49 | # You can specify the intermediary messages as tuples too. 50 | # ("human", "hi! my name is bob"), 51 | # ("ai", "Hello Bob! How can I assist you today?") 52 | HumanMessage(content="hi! my name is bob"), 53 | AIMessage(content="Hello Bob! How can I assist you today?"), 54 | ], 55 | "input": "whats the weather in sf?", 56 | }, 57 | config=RunnableConfig(run_name="Weather SF"), 58 | ) 59 | -------------------------------------------------------------------------------- /examples/langchain_variable.py: -------------------------------------------------------------------------------- 1 | from dotenv import load_dotenv 2 | from langchain.chat_models import init_chat_model 3 | 4 | from literalai import LiteralClient 5 | 6 | load_dotenv() 7 | 8 | lai = LiteralClient() 9 | lai.initialize() 10 | 11 | prompt = lai.api.get_or_create_prompt( 12 | name="user intent", 13 | template_messages=[ 14 | {"role": "system", "content": "You're a helpful assistant."}, 15 | {"role": "user", "content": "{{user_message}}"}, 16 | ], 17 | settings={ 18 | "provider": "openai", 19 | "model": "gpt-4o-mini", 20 | "temperature": 0, 21 | "max_tokens": 4095, 22 | "top_p": 1, 23 | "frequency_penalty": 0, 24 | "presence_penalty": 0, 25 | }, 26 | ) 27 | messages = prompt.to_langchain_chat_prompt_template() 28 | 29 | input_messages = messages.format_messages( 30 | user_message="The screen is cracked, there are scratches on the surface, and a component is missing." 31 | ) 32 | 33 | # Returns a langchain_openai.ChatOpenAI instance. 34 | gpt_4o = init_chat_model( # type: ignore 35 | model_provider=prompt.provider, 36 | **prompt.settings, 37 | ) 38 | 39 | lai.set_properties(prompt=prompt) 40 | print(gpt_4o.invoke(input_messages)) 41 | 42 | lai.flush_and_stop() 43 | -------------------------------------------------------------------------------- /examples/llamaindex.py: -------------------------------------------------------------------------------- 1 | from dotenv import load_dotenv 2 | from llama_index.core import Document, VectorStoreIndex 3 | 4 | from literalai import LiteralClient 5 | 6 | load_dotenv() 7 | 8 | client = LiteralClient() 9 | client.instrument_llamaindex() 10 | 11 | print("Vectorizing documents") 12 | index = VectorStoreIndex.from_documents([Document.example()]) 13 | query_engine = index.as_query_engine() 14 | 15 | questions = [ 16 | "Tell me about LLMs", 17 | "How do you fine-tune a neural network ?", 18 | "What is RAG ?", 19 | ] 20 | 21 | # No context, create a Thread (it will be named after the first user query) 22 | print(f"> \033[92m{questions[0]}\033[0m") 23 | response = query_engine.query(questions[0]) 24 | print(response) 25 | 26 | # Wrap in a thread (because it has no name it will also be named after the first user query) 27 | with client.thread() as thread: 28 | print(f"> \033[92m{questions[0]}\033[0m") 29 | response = query_engine.query(questions[0]) 30 | print(response) 31 | 32 | # Wrap in a thread (the name is conserved) 33 | with client.thread(name=f"User question : {questions[0]}") as thread: 34 | print(f"> \033[92m{questions[0]}\033[0m") 35 | response = query_engine.query(questions[0]) 36 | print(response) 37 | 38 | # One thread for all the questions 39 | with client.thread(name="Llamaindex questions") as thread: 40 | for question in questions: 41 | print(f"> \033[92m{question}\033[0m") 42 | response = query_engine.query(question) 43 | print(response) 44 | 45 | client.flush_and_stop() 46 | -------------------------------------------------------------------------------- /examples/llamaindex_workflow.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from llama_index.core.workflow import Event, StartEvent, StopEvent, Workflow, step 4 | from llama_index.llms.openai import OpenAI 5 | 6 | from literalai.client import LiteralClient 7 | 8 | lai_client = LiteralClient() 9 | lai_client.initialize() 10 | 11 | 12 | class JokeEvent(Event): 13 | joke: str 14 | 15 | 16 | class RewriteJoke(Event): 17 | joke: str 18 | 19 | 20 | class JokeFlow(Workflow): 21 | llm = OpenAI() 22 | 23 | @step() 24 | async def generate_joke(self, ev: StartEvent) -> JokeEvent: 25 | topic = ev.topic 26 | 27 | prompt = f"Write your best joke about {topic}." 28 | response = await self.llm.acomplete(prompt) 29 | return JokeEvent(joke=str(response)) 30 | 31 | @step() 32 | async def return_joke(self, ev: JokeEvent) -> RewriteJoke: 33 | return RewriteJoke(joke=ev.joke + "What is funny?") 34 | 35 | @step() 36 | async def critique_joke(self, ev: RewriteJoke) -> StopEvent: 37 | joke = ev.joke 38 | 39 | prompt = f"Give a thorough analysis and critique of the following joke: {joke}" 40 | response = await self.llm.acomplete(prompt) 41 | return StopEvent(result=str(response)) 42 | 43 | 44 | @lai_client.thread(name="JokeFlow") 45 | async def main(): 46 | w = JokeFlow(timeout=60, verbose=False) 47 | result = await w.run(topic="pirates") 48 | print(str(result)) 49 | 50 | 51 | if __name__ == "__main__": 52 | asyncio.run(main()) 53 | -------------------------------------------------------------------------------- /examples/main.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import json 3 | 4 | from dotenv import load_dotenv 5 | from openai import OpenAI 6 | 7 | from literalai import LiteralClient 8 | 9 | load_dotenv() 10 | 11 | client = OpenAI() 12 | 13 | sdk = LiteralClient(batch_size=2) 14 | sdk.instrument_openai() 15 | 16 | thread_id = None 17 | 18 | 19 | @sdk.step(type="run") 20 | def get_completion(welcome_message, text): 21 | completion = client.chat.completions.create( 22 | model="gpt-3.5-turbo", 23 | messages=[ 24 | { 25 | "role": "system", 26 | "content": "Tell an inspiring quote to the user, mentioning their name. Be extremely supportive while " 27 | "keeping it short. Write one sentence per line.", 28 | }, 29 | { 30 | "role": "assistant", 31 | "content": welcome_message, 32 | }, 33 | { 34 | "role": "user", 35 | "content": text, 36 | }, 37 | ], 38 | ) 39 | return completion.choices[0].message.content 40 | 41 | 42 | @sdk.thread 43 | def run(): 44 | global thread_id 45 | thread_id = sdk.get_current_thread().id 46 | 47 | welcome_message = "What's your name? " 48 | sdk.message(content=welcome_message, type="system_message") 49 | text = input(welcome_message) 50 | sdk.message(content=text, type="user_message") 51 | 52 | completion = get_completion(welcome_message=welcome_message, text=text) 53 | 54 | print("") 55 | print(completion) 56 | sdk.message(content=completion, type="assistant_message") 57 | 58 | 59 | run() 60 | sdk.flush_and_stop() 61 | 62 | 63 | # Get the steps from the API for the demo 64 | async def main(): 65 | print("\nSearching for the thread", thread_id, "...") 66 | thread = sdk.api.get_thread(id=thread_id) 67 | 68 | print(json.dumps(thread.to_dict(), indent=2)) 69 | 70 | # get the LLM step 71 | llm_step = [step for step in thread.steps if step.type == "llm"][0] 72 | 73 | if not llm_step: 74 | print("Error: No LLM step found") 75 | return 76 | 77 | # attach a score 78 | sdk.api.create_score( 79 | step_id=llm_step.id, 80 | name="user-feedback", 81 | type="HUMAN", 82 | value=1, 83 | comment="this is a comment", 84 | ) 85 | 86 | # get the updated steps 87 | thread = sdk.api.get_thread(id=thread_id) 88 | 89 | print( 90 | json.dumps( 91 | [step.to_dict() for step in thread.steps if step.type == "llm"], 92 | indent=2, 93 | ) 94 | ) 95 | 96 | 97 | asyncio.run(main()) 98 | -------------------------------------------------------------------------------- /examples/multimodal.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import time 3 | 4 | import requests # type: ignore 5 | from dotenv import load_dotenv 6 | from openai import OpenAI 7 | 8 | from literalai import LiteralClient 9 | 10 | load_dotenv() 11 | 12 | openai_client = OpenAI() 13 | 14 | literalai_client = LiteralClient() 15 | literalai_client.initialize() 16 | 17 | 18 | def encode_image(url): 19 | return base64.b64encode(requests.get(url).content) 20 | 21 | 22 | @literalai_client.step(type="run") 23 | def generate_answer(user_query, image_url): 24 | literalai_client.set_properties( 25 | name="foobar", 26 | metadata={"foo": "bar"}, 27 | tags=["foo", "bar"], 28 | ) 29 | completion = openai_client.chat.completions.create( 30 | model="gpt-4o-mini", 31 | messages=[ 32 | { 33 | "role": "user", 34 | "content": [ 35 | {"type": "text", "text": user_query}, 36 | { 37 | "type": "image_url", 38 | "image_url": {"url": image_url}, 39 | }, 40 | ], 41 | }, 42 | ], 43 | max_tokens=300, 44 | ) 45 | return completion.choices[0].message.content 46 | 47 | 48 | def main(): 49 | with literalai_client.thread(name="Meal Analyzer") as thread: 50 | welcome_message = ( 51 | "Welcome to the meal analyzer, please upload an image of your plate!" 52 | ) 53 | literalai_client.message( 54 | content=welcome_message, type="assistant_message", name="My Assistant" 55 | ) 56 | 57 | user_query = "Is this a healthy meal?" 58 | user_image = "https://www.eatthis.com/wp-content/uploads/sites/4/2021/05/healthy-plate.jpg" 59 | user_step = literalai_client.message( 60 | content=user_query, type="user_message", name="User" 61 | ) 62 | 63 | time.sleep(1) # to make sure the user step has arrived at Literal AI 64 | 65 | literalai_client.api.create_attachment( 66 | thread_id=thread.id, 67 | step_id=user_step.id, 68 | name="meal_image", 69 | content=encode_image(user_image), 70 | ) 71 | 72 | answer = generate_answer(user_query=user_query, image_url=user_image) 73 | literalai_client.message( 74 | content=answer, type="assistant_message", name="My Assistant" 75 | ) 76 | 77 | 78 | main() 79 | # Network requests by the SDK are performed asynchronously. 80 | # Invoke flush_and_stop() to guarantee the completion of all requests prior to the process termination. 81 | # WARNING: If you run a continuous server, you should not use this method. 82 | literalai_client.flush_and_stop() 83 | -------------------------------------------------------------------------------- /examples/openai_agents.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from agents import Agent, Runner, set_trace_processors, trace 4 | from dotenv import load_dotenv 5 | 6 | from literalai import LiteralClient 7 | 8 | load_dotenv() 9 | 10 | client = LiteralClient() 11 | 12 | 13 | async def main(): 14 | agent = Agent(name="Joke generator", instructions="Tell funny jokes.") 15 | 16 | with trace("Joke workflow"): 17 | first_result = await Runner.run(agent, "Tell me a joke") 18 | second_result = await Runner.run( 19 | agent, f"Rate this joke: {first_result.final_output}" 20 | ) 21 | print(f"Joke: {first_result.final_output}") 22 | print(f"Rating: {second_result.final_output}") 23 | 24 | 25 | if __name__ == "__main__": 26 | set_trace_processors([client.openai_agents_tracing_processor()]) 27 | asyncio.run(main()) 28 | -------------------------------------------------------------------------------- /examples/parallel_requests.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | 5 | from literalai import LiteralClient 6 | 7 | load_dotenv() 8 | 9 | client = LiteralClient() 10 | 11 | 12 | @client.thread 13 | def request(query): 14 | client.message(query) 15 | return query 16 | 17 | 18 | @client.thread 19 | async def async_request(query, sleepy): 20 | client.message(query) 21 | await asyncio.sleep(sleepy) 22 | return query 23 | 24 | 25 | @client.thread(thread_id="e3fcf535-2555-4f75-bc10-fc1499baeff4") 26 | def precise_request(query): 27 | client.message(query) 28 | return query 29 | 30 | 31 | async def main(): 32 | request("hello") 33 | request("world!") 34 | precise_request("bonjour!") 35 | 36 | await asyncio.gather(async_request("foo", 5), async_request("bar", 2)) 37 | 38 | 39 | asyncio.run(main()) 40 | client.flush_and_stop() 41 | -------------------------------------------------------------------------------- /examples/prompt.py: -------------------------------------------------------------------------------- 1 | from dotenv import load_dotenv 2 | 3 | from literalai import LiteralClient 4 | 5 | load_dotenv() 6 | 7 | client = LiteralClient() 8 | 9 | prompt = client.api.get_prompt(name="Default", version=0) 10 | 11 | print(prompt) 12 | -------------------------------------------------------------------------------- /examples/samplefile.txt: -------------------------------------------------------------------------------- 1 | foobar 2 | -------------------------------------------------------------------------------- /examples/streaming.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | from openai import AsyncOpenAI, OpenAI 5 | 6 | from literalai import LiteralClient 7 | 8 | load_dotenv() 9 | 10 | async_client = AsyncOpenAI() 11 | client = OpenAI() 12 | 13 | 14 | sdk = LiteralClient(batch_size=2) 15 | sdk.initialize() 16 | 17 | 18 | @sdk.thread 19 | async def async_run(): 20 | with sdk.step(type="run", name="async_get_chat_completion"): 21 | stream = await async_client.chat.completions.create( 22 | model="gpt-3.5-turbo", 23 | messages=[ 24 | { 25 | "role": "system", 26 | "content": "Tell an inspiring quote to the user, mentioning their name. Be extremely supportive while keeping it short. Write one sentence per line.", 27 | }, 28 | { 29 | "role": "assistant", 30 | "content": "What's your name? ", 31 | }, 32 | { 33 | "role": "user", 34 | "content": "Joe", 35 | }, 36 | ], 37 | stream=True, 38 | ) 39 | async for chunk in stream: 40 | if chunk.choices[0].delta.content is not None: 41 | print(chunk.choices[0].delta.content, end="") 42 | print("") 43 | 44 | with sdk.step(type="run", name="async_get_completion"): 45 | stream = await async_client.completions.create( 46 | model="gpt-3.5-turbo-instruct", 47 | prompt="Tell an inspiring quote to the user, mentioning their name. Be extremely supportive while keeping it short. Write one sentence per line.\n\nAssistant: What's your name?\n\nUser: Joe\n\nAssistant: ", 48 | stream=True, 49 | ) 50 | async for chunk in stream: 51 | if chunk.choices[0].text is not None: 52 | print(chunk.choices[0].text, end="") 53 | print("") 54 | 55 | 56 | asyncio.run(async_run()) 57 | 58 | 59 | @sdk.thread 60 | def run(): 61 | with sdk.step(type="run", name="get_chat_completion"): 62 | stream = client.chat.completions.create( 63 | model="gpt-3.5-turbo", 64 | messages=[ 65 | { 66 | "role": "system", 67 | "content": "Tell an inspiring quote to the user, mentioning their name. Be extremely supportive while keeping it short. Write one sentence per line.", 68 | }, 69 | { 70 | "role": "assistant", 71 | "content": "What's your name? ", 72 | }, 73 | { 74 | "role": "user", 75 | "content": "Joe", 76 | }, 77 | ], 78 | stream=True, 79 | ) 80 | for chunk in stream: 81 | if chunk.choices[0].delta.content is not None: 82 | print(chunk.choices[0].delta.content, end="") 83 | print("") 84 | 85 | with sdk.step(type="run", name="get_completion"): 86 | stream = client.completions.create( 87 | model="gpt-3.5-turbo-instruct", 88 | prompt="Tell an inspiring quote to the user, mentioning their name. Be extremely supportive while keeping it short. Write one sentence per line.\n\nAssistant: What's your name?\n\nUser: Joe\n\nAssistant: ", 89 | stream=True, 90 | ) 91 | for chunk in stream: 92 | if chunk.choices[0].text is not None: 93 | print(chunk.choices[0].text, end="") 94 | print("") 95 | 96 | 97 | run() 98 | sdk.flush_and_stop() 99 | -------------------------------------------------------------------------------- /examples/thread.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | 5 | from literalai import LiteralClient 6 | 7 | load_dotenv() 8 | 9 | 10 | sdk = LiteralClient(batch_size=2) 11 | 12 | 13 | async def main(): 14 | thread = sdk.api.create_thread(metadata={"key": "value"}, tags=["hello"]) 15 | 16 | id = thread.id 17 | print(id, thread.to_dict()) 18 | 19 | thread = sdk.api.update_thread( 20 | id=id, 21 | metadata={"test": "test"}, 22 | tags=["hello:world"], 23 | ) 24 | 25 | print(thread.to_dict()) 26 | 27 | thread = sdk.api.get_thread(id=id) 28 | 29 | print(thread.to_dict()) 30 | 31 | sdk.api.delete_thread(id=id) 32 | 33 | try: 34 | thread = sdk.api.get_thread(id=id) 35 | except Exception as e: 36 | print(e) 37 | 38 | after = None 39 | max_calls = 5 40 | while len((result := sdk.api.get_threads(first=2, after=after)).data) > 0: 41 | print(result.to_dict()) 42 | after = result.pageInfo.endCursor 43 | max_calls -= 1 44 | if max_calls == 0: 45 | break 46 | 47 | print("filtered") 48 | 49 | threads = sdk.api.get_threads( 50 | filters=[ 51 | { 52 | "field": "createdAt", 53 | "operator": "gt", 54 | "value": "2023-12-05", 55 | }, 56 | ], 57 | # order_by={"column": "participant", "direction": "ASC"}, 58 | ) 59 | print(threads.to_dict()) 60 | 61 | 62 | asyncio.run(main()) 63 | -------------------------------------------------------------------------------- /examples/user.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | from dotenv import load_dotenv 4 | 5 | from literalai import LiteralClient 6 | 7 | load_dotenv() 8 | 9 | 10 | sdk = LiteralClient(batch_size=2) 11 | 12 | 13 | async def main(): 14 | user = sdk.api.create_user(identifier="test-user-example", metadata={"name": "123"}) 15 | 16 | id = user.id 17 | print(id, user.to_dict()) 18 | 19 | user = sdk.api.update_user( 20 | id=id, 21 | identifier="user", 22 | metadata={"test": "test"}, 23 | ) 24 | 25 | print(user.to_dict()) 26 | 27 | user = sdk.api.get_user(id=id) 28 | 29 | print(user.to_dict()) 30 | 31 | user = sdk.api.get_user(identifier="user") 32 | 33 | print(user.to_dict()) 34 | 35 | sdk.api.delete_user(id=id) 36 | 37 | try: 38 | user = sdk.api.update_user( 39 | id=id, 40 | identifier="user", 41 | metadata={"test": "test"}, 42 | ) 43 | except Exception as e: 44 | print(e) 45 | 46 | 47 | asyncio.run(main()) 48 | -------------------------------------------------------------------------------- /literalai/__init__.py: -------------------------------------------------------------------------------- 1 | from literalai.client import AsyncLiteralClient, LiteralClient 2 | from literalai.evaluation.dataset import Dataset 3 | from literalai.evaluation.dataset_experiment import ( 4 | DatasetExperiment, 5 | DatasetExperimentItem, 6 | ) 7 | from literalai.evaluation.dataset_item import DatasetItem 8 | from literalai.my_types import * # noqa 9 | from literalai.observability.generation import ( 10 | BaseGeneration, 11 | ChatGeneration, 12 | CompletionGeneration, 13 | GenerationMessage, 14 | ) 15 | from literalai.observability.message import Message 16 | from literalai.observability.step import Attachment, Score, Step 17 | from literalai.observability.thread import Thread 18 | from literalai.prompt_engineering.prompt import Prompt 19 | from literalai.version import __version__ 20 | 21 | __all__ = [ 22 | "LiteralClient", 23 | "AsyncLiteralClient", 24 | "BaseGeneration", 25 | "CompletionGeneration", 26 | "ChatGeneration", 27 | "GenerationMessage", 28 | "Message", 29 | "Step", 30 | "Score", 31 | "Thread", 32 | "Dataset", 33 | "Attachment", 34 | "DatasetItem", 35 | "DatasetExperiment", 36 | "DatasetExperimentItem", 37 | "Prompt", 38 | "__version__", 39 | ] 40 | -------------------------------------------------------------------------------- /literalai/api/README.md: -------------------------------------------------------------------------------- 1 | # Literal AI API 2 | 3 | This module contains the APIs to directly interact with the Literal AI platform. 4 | 5 | `BaseLiteralAPI` has all the methods prototypes and is the source of truth when it comes to documentation. 6 | 7 | Inheriting from `BaseLiteralAPI` are `LiteralAPI` and `AsyncLiteralAPI`, the synchronous and asynchronous APIs respectively. 8 | -------------------------------------------------------------------------------- /literalai/api/__init__.py: -------------------------------------------------------------------------------- 1 | from literalai.api.asynchronous import AsyncLiteralAPI 2 | from literalai.api.synchronous import LiteralAPI 3 | 4 | __all__ = ["LiteralAPI", "AsyncLiteralAPI"] 5 | -------------------------------------------------------------------------------- /literalai/api/helpers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/api/helpers/__init__.py -------------------------------------------------------------------------------- /literalai/api/helpers/attachment_helpers.py: -------------------------------------------------------------------------------- 1 | import mimetypes 2 | from typing import Dict, Optional, TypedDict, Union 3 | 4 | from literalai.api.helpers import gql 5 | from literalai.observability.step import Attachment 6 | 7 | 8 | def create_attachment_helper( 9 | step_id: str, 10 | thread_id: Optional[str] = None, 11 | id: Optional[str] = None, 12 | metadata: Optional[Dict] = None, 13 | mime: Optional[str] = None, 14 | name: Optional[str] = None, 15 | object_key: Optional[str] = None, 16 | url: Optional[str] = None, 17 | content: Optional[Union[bytes, str]] = None, 18 | path: Optional[str] = None, 19 | ): 20 | if not content and not url and not path: 21 | raise Exception("Either content, path or attachment url must be provided") 22 | 23 | if content and path: 24 | raise Exception("Only one of content and path must be provided") 25 | 26 | if (content and url) or (path and url): 27 | raise Exception("Only one of content, path and attachment url must be provided") 28 | 29 | if path: 30 | # TODO: if attachment.mime is text, we could read as text? 31 | with open(path, "rb") as f: 32 | content = f.read() 33 | if not name: 34 | name = path.split("/")[-1] 35 | if not mime: 36 | mime, _ = mimetypes.guess_type(path) 37 | mime = mime or "application/octet-stream" 38 | 39 | if not name: 40 | raise Exception("Attachment name must be provided") 41 | 42 | variables = { 43 | "metadata": metadata, 44 | "mime": mime, 45 | "name": name, 46 | "objectKey": object_key, 47 | "stepId": step_id, 48 | "threadId": thread_id, 49 | "url": url, 50 | "id": id, 51 | } 52 | 53 | description = "create attachment" 54 | 55 | def process_response(response): 56 | return Attachment.from_dict(response["data"]["createAttachment"]) 57 | 58 | return gql.CREATE_ATTACHMENT, description, variables, content, process_response 59 | 60 | 61 | class AttachmentUpload(TypedDict, total=False): 62 | metadata: Optional[Dict] 63 | name: Optional[str] 64 | mime: Optional[str] 65 | objectKey: Optional[str] 66 | url: Optional[str] 67 | 68 | 69 | def update_attachment_helper(id: str, update_params: AttachmentUpload): 70 | variables = {"id": id, **update_params} 71 | 72 | def process_response(response): 73 | return Attachment.from_dict(response["data"]["updateAttachment"]) 74 | 75 | description = "update attachment" 76 | 77 | return gql.UPDATE_ATTACHMENT, description, variables, process_response 78 | 79 | 80 | def get_attachment_helper(id: str): 81 | variables = {"id": id} 82 | 83 | def process_response(response): 84 | attachment = response["data"]["attachment"] 85 | return Attachment.from_dict(attachment) if attachment else None 86 | 87 | description = "get attachment" 88 | 89 | return gql.GET_ATTACHMENT, description, variables, process_response 90 | 91 | 92 | def delete_attachment_helper(id: str): 93 | variables = {"id": id} 94 | 95 | def process_response(response): 96 | return response["data"]["deleteAttachment"] 97 | 98 | description = "delete attachment" 99 | 100 | return gql.DELETE_ATTACHMENT, description, variables, process_response 101 | -------------------------------------------------------------------------------- /literalai/api/helpers/dataset_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import TYPE_CHECKING, Dict, Optional 2 | 3 | from literalai.api.helpers import gql 4 | from literalai.evaluation.dataset import Dataset, DatasetType 5 | from literalai.evaluation.dataset_experiment import ( 6 | DatasetExperiment, 7 | DatasetExperimentItem, 8 | ) 9 | from literalai.evaluation.dataset_item import DatasetItem 10 | 11 | if TYPE_CHECKING: 12 | from literalai.api import LiteralAPI 13 | 14 | 15 | def create_dataset_helper( 16 | api: "LiteralAPI", 17 | name: str, 18 | description: Optional[str] = None, 19 | metadata: Optional[Dict] = None, 20 | type: DatasetType = "key_value", 21 | ): 22 | variables = { 23 | "name": name, 24 | "description": description, 25 | "metadata": metadata, 26 | "type": type, 27 | } 28 | 29 | def process_response(response): 30 | return Dataset.from_dict(api, response["data"]["createDataset"]) 31 | 32 | description = "create dataset" 33 | 34 | return gql.CREATE_DATASET, description, variables, process_response 35 | 36 | 37 | def get_dataset_helper( 38 | api: "LiteralAPI", id: Optional[str] = None, name: Optional[str] = None 39 | ): 40 | if not id and not name: 41 | raise ValueError("id or name must be provided") 42 | 43 | body = {} 44 | 45 | if id: 46 | body["id"] = id 47 | if name: 48 | body["name"] = name 49 | 50 | def process_response(response) -> Optional[Dataset]: 51 | dataset_dict = response.get("data") 52 | if dataset_dict is None: 53 | return None 54 | return Dataset.from_dict(api, dataset_dict) 55 | 56 | description = "get dataset" 57 | 58 | # Assuming there's a placeholder or a constant for the subpath 59 | subpath = "/export/dataset" 60 | 61 | return subpath, description, body, process_response 62 | 63 | 64 | def update_dataset_helper( 65 | api: "LiteralAPI", 66 | id: str, 67 | name: Optional[str] = None, 68 | description: Optional[str] = None, 69 | metadata: Optional[Dict] = None, 70 | ): 71 | variables: Dict = {"id": id} 72 | if name is not None: 73 | variables["name"] = name 74 | if description is not None: 75 | variables["description"] = description 76 | if metadata is not None: 77 | variables["metadata"] = metadata 78 | 79 | def process_response(response): 80 | return Dataset.from_dict(api, response["data"]["updateDataset"]) 81 | 82 | description = "update dataset" 83 | 84 | return gql.UPDATE_DATASET, description, variables, process_response 85 | 86 | 87 | def delete_dataset_helper(api: "LiteralAPI", id: str): 88 | variables = {"id": id} 89 | 90 | def process_response(response): 91 | return Dataset.from_dict(api, response["data"]["deleteDataset"]) 92 | 93 | description = "delete dataset" 94 | 95 | return gql.DELETE_DATASET, description, variables, process_response 96 | 97 | 98 | def create_experiment_helper( 99 | api: "LiteralAPI", 100 | name: str, 101 | dataset_id: Optional[str] = None, 102 | prompt_variant_id: Optional[str] = None, 103 | params: Optional[Dict] = None, 104 | ): 105 | variables = { 106 | "datasetId": dataset_id, 107 | "name": name, 108 | "promptExperimentId": prompt_variant_id, 109 | "params": params, 110 | } 111 | 112 | def process_response(response): 113 | return DatasetExperiment.from_dict( 114 | api, response["data"]["createDatasetExperiment"] 115 | ) 116 | 117 | description = "create dataset experiment" 118 | 119 | return gql.CREATE_EXPERIMENT, description, variables, process_response 120 | 121 | 122 | def create_experiment_item_helper( 123 | dataset_experiment_id: str, 124 | experiment_run_id: Optional[str] = None, 125 | dataset_item_id: Optional[str] = None, 126 | input: Optional[Dict] = None, 127 | output: Optional[Dict] = None, 128 | ): 129 | variables = { 130 | "experimentRunId": experiment_run_id, 131 | "datasetExperimentId": dataset_experiment_id, 132 | "datasetItemId": dataset_item_id, 133 | "input": input, 134 | "output": output, 135 | } 136 | 137 | def process_response(response): 138 | return DatasetExperimentItem.from_dict( 139 | response["data"]["createDatasetExperimentItem"] 140 | ) 141 | 142 | description = "create dataset experiment item" 143 | 144 | return gql.CREATE_EXPERIMENT_ITEM, description, variables, process_response 145 | 146 | 147 | def create_dataset_item_helper( 148 | dataset_id: str, 149 | input: Dict, 150 | expected_output: Optional[Dict] = None, 151 | metadata: Optional[Dict] = None, 152 | ): 153 | variables = { 154 | "datasetId": dataset_id, 155 | "input": input, 156 | "expectedOutput": expected_output, 157 | "metadata": metadata, 158 | } 159 | 160 | def process_response(response): 161 | return DatasetItem.from_dict(response["data"]["createDatasetItem"]) 162 | 163 | description = "create dataset item" 164 | 165 | return gql.CREATE_DATASET_ITEM, description, variables, process_response 166 | 167 | 168 | def get_dataset_item_helper(id: str): 169 | variables = {"id": id} 170 | 171 | def process_response(response): 172 | return DatasetItem.from_dict(response["data"]["datasetItem"]) 173 | 174 | description = "get dataset item" 175 | 176 | return gql.GET_DATASET_ITEM, description, variables, process_response 177 | 178 | 179 | def delete_dataset_item_helper(id: str): 180 | variables = {"id": id} 181 | 182 | def process_response(response): 183 | return DatasetItem.from_dict(response["data"]["deleteDatasetItem"]) 184 | 185 | description = "delete dataset item" 186 | 187 | return gql.DELETE_DATASET_ITEM, description, variables, process_response 188 | 189 | 190 | def add_step_to_dataset_helper( 191 | dataset_id: str, step_id: str, metadata: Optional[Dict] = None 192 | ): 193 | variables = { 194 | "datasetId": dataset_id, 195 | "stepId": step_id, 196 | "metadata": metadata, 197 | } 198 | 199 | def process_response(response): 200 | return DatasetItem.from_dict(response["data"]["addStepToDataset"]) 201 | 202 | description = "add step to dataset" 203 | 204 | return gql.ADD_STEP_TO_DATASET, description, variables, process_response 205 | 206 | 207 | def add_generation_to_dataset_helper( 208 | dataset_id: str, generation_id: str, metadata: Optional[Dict] = None 209 | ): 210 | variables = { 211 | "datasetId": dataset_id, 212 | "generationId": generation_id, 213 | "metadata": metadata, 214 | } 215 | 216 | def process_response(response): 217 | return DatasetItem.from_dict(response["data"]["addGenerationToDataset"]) 218 | 219 | description = "add generation to dataset" 220 | 221 | return gql.ADD_GENERATION_TO_DATASET, description, variables, process_response 222 | -------------------------------------------------------------------------------- /literalai/api/helpers/generation_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional, Union 2 | 3 | from literalai.api.helpers import gql 4 | from literalai.my_types import PaginatedResponse 5 | from literalai.observability.filter import generations_filters, generations_order_by 6 | from literalai.observability.generation import ( 7 | BaseGeneration, 8 | ChatGeneration, 9 | CompletionGeneration, 10 | ) 11 | 12 | 13 | def get_generations_helper( 14 | first: Optional[int] = None, 15 | after: Optional[str] = None, 16 | before: Optional[str] = None, 17 | filters: Optional[generations_filters] = None, 18 | order_by: Optional[generations_order_by] = None, 19 | ): 20 | variables: Dict[str, Any] = {} 21 | 22 | if first: 23 | variables["first"] = first 24 | if after: 25 | variables["after"] = after 26 | if before: 27 | variables["before"] = before 28 | if filters: 29 | variables["filters"] = filters 30 | if order_by: 31 | variables["orderBy"] = order_by 32 | 33 | def process_response(response): 34 | processed_response = response["data"]["generations"] 35 | processed_response["data"] = list( 36 | map(lambda x: x["node"], processed_response["edges"]) 37 | ) 38 | return PaginatedResponse[BaseGeneration].from_dict( 39 | processed_response, BaseGeneration 40 | ) 41 | 42 | description = "get generations" 43 | 44 | return gql.GET_GENERATIONS, description, variables, process_response 45 | 46 | 47 | def create_generation_helper(generation: Union[ChatGeneration, CompletionGeneration]): 48 | variables = {"generation": generation.to_dict()} 49 | 50 | def process_response(response): 51 | return BaseGeneration.from_dict(response["data"]["createGeneration"]) 52 | 53 | description = "create generation" 54 | 55 | return gql.CREATE_GENERATION, description, variables, process_response 56 | -------------------------------------------------------------------------------- /literalai/api/helpers/prompt_helpers.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import TYPE_CHECKING, Callable, Optional, TypedDict 3 | 4 | from literalai.cache.prompt_helpers import put_prompt 5 | from literalai.observability.generation import GenerationMessage 6 | from literalai.prompt_engineering.prompt import Prompt, ProviderSettings 7 | 8 | if TYPE_CHECKING: 9 | from literalai.api import LiteralAPI 10 | from literalai.cache.shared_cache import SharedCache 11 | 12 | from literalai.api.helpers import gql 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | 17 | def create_prompt_lineage_helper(name: str, description: Optional[str] = None): 18 | variables = {"name": name, "description": description} 19 | 20 | def process_response(response): 21 | prompt = response["data"]["createPromptLineage"] 22 | if prompt and prompt.get("deletedAt"): 23 | logger.warning( 24 | f"Prompt {name} was deleted - please update any references to use an active prompt in production" 25 | ) 26 | return prompt 27 | 28 | description = "create prompt lineage" 29 | 30 | return gql.CREATE_PROMPT_LINEAGE, description, variables, process_response 31 | 32 | 33 | def get_prompt_lineage_helper(name: str): 34 | variables = {"name": name} 35 | 36 | def process_response(response): 37 | prompt = response["data"]["promptLineage"] 38 | if prompt and prompt.get("deletedAt"): 39 | logger.warning( 40 | f"Prompt {name} was deleted - please update any references to use an active prompt in production" 41 | ) 42 | return prompt 43 | 44 | description = "get prompt lineage" 45 | 46 | return gql.GET_PROMPT_LINEAGE, description, variables, process_response 47 | 48 | 49 | def create_prompt_helper( 50 | api: "LiteralAPI", 51 | lineage_id: str, 52 | template_messages: list[GenerationMessage], 53 | settings: Optional[ProviderSettings] = None, 54 | tools: Optional[list[dict]] = None, 55 | ): 56 | variables = { 57 | "lineageId": lineage_id, 58 | "templateMessages": template_messages, 59 | "settings": settings, 60 | "tools": tools, 61 | } 62 | 63 | def process_response(response): 64 | prompt = response["data"]["createPromptVersion"] 65 | prompt_lineage = prompt.get("lineage") 66 | 67 | if prompt_lineage and prompt_lineage.get("deletedAt"): 68 | logger.warning( 69 | f"Prompt {prompt_lineage.get('name')} was deleted - please update any references to use an active prompt in production" 70 | ) 71 | return Prompt.from_dict(api, prompt) if prompt else None 72 | 73 | description = "create prompt version" 74 | 75 | return gql.CREATE_PROMPT_VERSION, description, variables, process_response 76 | 77 | 78 | def get_prompt_cache_key( 79 | id: Optional[str], name: Optional[str], version: Optional[int] 80 | ) -> str: 81 | if id: 82 | return id 83 | elif name and version: 84 | return f"{name}-{version}" 85 | elif name: 86 | return name 87 | else: 88 | raise ValueError("Either the `id` or the `name` must be provided.") 89 | 90 | 91 | def get_prompt_helper( 92 | api: "LiteralAPI", 93 | id: Optional[str] = None, 94 | name: Optional[str] = None, 95 | version: Optional[int] = 0, 96 | cache: Optional["SharedCache"] = None, 97 | ) -> tuple[str, str, dict, Callable, int, Optional[Prompt]]: 98 | """Helper function for getting prompts with caching logic""" 99 | 100 | cached_prompt = None 101 | timeout = 10 102 | 103 | if cache: 104 | cached_prompt = cache.get(get_prompt_cache_key(id, name, version)) 105 | timeout = 1 if cached_prompt else timeout 106 | 107 | variables = {"id": id, "name": name, "version": version} 108 | 109 | def process_response(response): 110 | prompt_version = response["data"]["promptVersion"] 111 | prompt_lineage = prompt_version.get("lineage") 112 | 113 | if prompt_lineage and prompt_lineage.get("deletedAt"): 114 | logger.warning( 115 | f"Prompt {name} was deleted - please update any references to use an active prompt in production" 116 | ) 117 | prompt = Prompt.from_dict(api, prompt_version) if prompt_version else None 118 | if cache and prompt: 119 | put_prompt(cache, prompt) 120 | return prompt 121 | 122 | description = "get prompt" 123 | 124 | return ( 125 | gql.GET_PROMPT_VERSION, 126 | description, 127 | variables, 128 | process_response, 129 | timeout, 130 | cached_prompt, 131 | ) 132 | 133 | 134 | def create_prompt_variant_helper( 135 | from_lineage_id: Optional[str] = None, 136 | template_messages: list[GenerationMessage] = [], 137 | settings: Optional[ProviderSettings] = None, 138 | tools: Optional[list[dict]] = None, 139 | ): 140 | variables = { 141 | "fromLineageId": from_lineage_id, 142 | "templateMessages": template_messages, 143 | "settings": settings, 144 | "tools": tools, 145 | } 146 | 147 | def process_response(response) -> Optional[str]: 148 | variant = response["data"]["createPromptExperiment"] 149 | return variant["id"] if variant else None 150 | 151 | description = "create prompt variant" 152 | 153 | return gql.CREATE_PROMPT_VARIANT, description, variables, process_response 154 | 155 | 156 | class PromptRollout(TypedDict): 157 | version: int 158 | rollout: int 159 | 160 | 161 | def get_prompt_ab_testing_helper( 162 | name: Optional[str] = None, 163 | ): 164 | variables = {"lineageName": name} 165 | 166 | def process_response(response) -> list[PromptRollout]: 167 | response_data = response["data"]["promptLineageRollout"] 168 | return list(map(lambda x: x["node"], response_data["edges"])) 169 | 170 | description = "get prompt A/B testing" 171 | 172 | return gql.GET_PROMPT_AB_TESTING, description, variables, process_response 173 | 174 | 175 | def update_prompt_ab_testing_helper(name: str, rollouts: list[PromptRollout]): 176 | variables = {"name": name, "rollouts": rollouts} 177 | 178 | def process_response(response) -> dict: 179 | return response["data"]["updatePromptLineageRollout"] 180 | 181 | description = "update prompt A/B testing" 182 | 183 | return gql.UPDATE_PROMPT_AB_TESTING, description, variables, process_response 184 | -------------------------------------------------------------------------------- /literalai/api/helpers/score_helpers.py: -------------------------------------------------------------------------------- 1 | import math 2 | from typing import Any, Dict, List, Optional, TypedDict 3 | 4 | from literalai.api.helpers import gql 5 | from literalai.my_types import PaginatedResponse 6 | from literalai.observability.filter import scores_filters, scores_order_by 7 | from literalai.observability.step import Score, ScoreDict, ScoreType 8 | 9 | 10 | def get_scores_helper( 11 | first: Optional[int] = None, 12 | after: Optional[str] = None, 13 | before: Optional[str] = None, 14 | filters: Optional[scores_filters] = None, 15 | order_by: Optional[scores_order_by] = None, 16 | ): 17 | variables: Dict[str, Any] = {} 18 | 19 | if first: 20 | variables["first"] = first 21 | if after: 22 | variables["after"] = after 23 | if before: 24 | variables["before"] = before 25 | if filters: 26 | variables["filters"] = filters 27 | if order_by: 28 | variables["orderBy"] = order_by 29 | 30 | def process_response(response): 31 | response_data = response["data"]["scores"] 32 | response_data["data"] = list(map(lambda x: x["node"], response_data["edges"])) 33 | return PaginatedResponse[Score].from_dict(response_data, Score) # type: ignore 34 | 35 | description = "get scores" 36 | 37 | return gql.GET_SCORES, description, variables, process_response 38 | 39 | 40 | def create_score_helper( 41 | name: str, 42 | value: float, 43 | type: ScoreType, 44 | step_id: Optional[str] = None, 45 | dataset_experiment_item_id: Optional[str] = None, 46 | comment: Optional[str] = None, 47 | tags: Optional[List[str]] = None, 48 | ): 49 | variables = { 50 | "name": name, 51 | "type": type, 52 | "value": value, 53 | "stepId": step_id, 54 | "datasetExperimentItemId": dataset_experiment_item_id, 55 | "comment": comment, 56 | "tags": tags, 57 | } 58 | 59 | def process_response(response): 60 | return Score.from_dict(response["data"]["createScore"]) 61 | 62 | description = "create score" 63 | 64 | return gql.CREATE_SCORE, description, variables, process_response 65 | 66 | 67 | def create_scores_fields_builder(scores: List[ScoreDict]): 68 | generated = "" 69 | for id in range(len(scores)): 70 | generated += f"""$name_{id}: String! 71 | $type_{id}: ScoreType! 72 | $value_{id}: Float! 73 | $label_{id}: String 74 | $stepId_{id}: String 75 | $datasetExperimentItemId_{id}: String 76 | $scorer_{id}: String 77 | $comment_{id}: String 78 | $tags_{id}: [String!] 79 | """ 80 | return generated 81 | 82 | 83 | def create_scores_args_builder(scores: List[ScoreDict]): 84 | generated = "" 85 | for id in range(len(scores)): 86 | generated += f""" 87 | score_{id}: createScore( 88 | name: $name_{id} 89 | type: $type_{id} 90 | value: $value_{id} 91 | label: $label_{id} 92 | stepId: $stepId_{id} 93 | datasetExperimentItemId: $datasetExperimentItemId_{id} 94 | scorer: $scorer_{id} 95 | comment: $comment_{id} 96 | tags: $tags_{id} 97 | ) {{ 98 | id 99 | name 100 | type 101 | value 102 | label 103 | comment 104 | scorer 105 | }} 106 | """ 107 | return generated 108 | 109 | 110 | def create_scores_query_builder(scores: List[ScoreDict]): 111 | return f""" 112 | mutation CreateScores({create_scores_fields_builder(scores)}) {{ 113 | {create_scores_args_builder(scores)} 114 | }} 115 | """ 116 | 117 | 118 | class ScoreUpdate(TypedDict, total=False): 119 | comment: Optional[str] 120 | value: float 121 | 122 | 123 | def update_score_helper( 124 | id: str, 125 | update_params: ScoreUpdate, 126 | ): 127 | variables = {"id": id, **update_params} 128 | 129 | def process_response(response): 130 | return Score.from_dict(response["data"]["updateScore"]) 131 | 132 | description = "update score" 133 | 134 | return gql.UPDATE_SCORE, description, variables, process_response 135 | 136 | 137 | def delete_score_helper(id: str): 138 | variables = {"id": id} 139 | 140 | def process_response(response): 141 | return response["data"]["deleteScore"] 142 | 143 | description = "delete score" 144 | 145 | return gql.DELETE_SCORE, description, variables, process_response 146 | 147 | 148 | def check_scores_finite(scores: List[ScoreDict]): 149 | for score in scores: 150 | if not math.isfinite(score["value"]): 151 | raise ValueError( 152 | f"Value {score['value']} for score {score['name']} is not finite" 153 | ) 154 | return True 155 | -------------------------------------------------------------------------------- /literalai/api/helpers/step_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Optional, Union 2 | 3 | from literalai.api.helpers import gql 4 | from literalai.my_types import PaginatedResponse 5 | from literalai.observability.filter import steps_filters, steps_order_by 6 | from literalai.observability.step import Step, StepDict, StepType 7 | 8 | 9 | def create_step_helper( 10 | thread_id: Optional[str] = None, 11 | type: Optional[StepType] = "undefined", 12 | start_time: Optional[str] = None, 13 | end_time: Optional[str] = None, 14 | input: Optional[Dict] = None, 15 | output: Optional[Dict] = None, 16 | metadata: Optional[Dict] = None, 17 | parent_id: Optional[str] = None, 18 | name: Optional[str] = None, 19 | tags: Optional[List[str]] = None, 20 | root_run_id: Optional[str] = None, 21 | ): 22 | variables = { 23 | "threadId": thread_id, 24 | "type": type, 25 | "startTime": start_time, 26 | "endTime": end_time, 27 | "input": input, 28 | "output": output, 29 | "metadata": metadata, 30 | "parentId": parent_id, 31 | "name": name, 32 | "tags": tags, 33 | "root_run_id": root_run_id, 34 | } 35 | 36 | def process_response(response): 37 | return Step.from_dict(response["data"]["createStep"]) 38 | 39 | description = "create step" 40 | 41 | return gql.CREATE_STEP, description, variables, process_response 42 | 43 | 44 | def update_step_helper( 45 | id: str, 46 | type: Optional[StepType] = None, 47 | input: Optional[str] = None, 48 | output: Optional[str] = None, 49 | metadata: Optional[Dict] = None, 50 | name: Optional[str] = None, 51 | tags: Optional[List[str]] = None, 52 | start_time: Optional[str] = None, 53 | end_time: Optional[str] = None, 54 | parent_id: Optional[str] = None, 55 | ): 56 | variables = { 57 | "id": id, 58 | "type": type, 59 | "input": input, 60 | "output": output, 61 | "metadata": metadata, 62 | "name": name, 63 | "tags": tags, 64 | "startTime": start_time, 65 | "endTime": end_time, 66 | "parentId": parent_id, 67 | } 68 | 69 | def process_response(response): 70 | return Step.from_dict(response["data"]["updateStep"]) 71 | 72 | description = "update step" 73 | 74 | return gql.UPDATE_STEP, description, variables, process_response 75 | 76 | 77 | def get_steps_helper( 78 | first: Optional[int] = None, 79 | after: Optional[str] = None, 80 | before: Optional[str] = None, 81 | filters: Optional[steps_filters] = None, 82 | order_by: Optional[steps_order_by] = None, 83 | ): 84 | variables: Dict[str, Any] = {} 85 | 86 | if first: 87 | variables["first"] = first 88 | if after: 89 | variables["after"] = after 90 | if before: 91 | variables["before"] = before 92 | if filters: 93 | variables["filters"] = filters 94 | if order_by: 95 | variables["orderBy"] = order_by 96 | 97 | def process_response(response): 98 | processed_response = response["data"]["steps"] 99 | processed_response["data"] = [ 100 | edge["node"] for edge in processed_response["edges"] 101 | ] 102 | return PaginatedResponse[Step].from_dict(processed_response, Step) 103 | 104 | description = "get steps" 105 | 106 | return gql.GET_STEPS, description, variables, process_response 107 | 108 | 109 | def get_step_helper(id: str): 110 | variables = {"id": id} 111 | 112 | def process_response(response): 113 | step = response["data"]["step"] 114 | return Step.from_dict(step) if step else None 115 | 116 | description = "get step" 117 | 118 | return gql.GET_STEP, description, variables, process_response 119 | 120 | 121 | def delete_step_helper(id: str): 122 | variables = {"id": id} 123 | 124 | def process_response(response): 125 | return bool(response["data"]["deleteStep"]) 126 | 127 | description = "delete step" 128 | 129 | return gql.DELETE_STEP, description, variables, process_response 130 | 131 | 132 | def send_steps_helper(steps: List[Union[StepDict, "Step"]]): 133 | query = gql.steps_query_builder(steps) 134 | variables = gql.steps_variables_builder(steps) 135 | 136 | description = "send steps" 137 | 138 | def process_response(response: Dict): 139 | return response 140 | 141 | return query, description, variables, process_response 142 | -------------------------------------------------------------------------------- /literalai/api/helpers/thread_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Optional 2 | 3 | from literalai.api.helpers import gql 4 | from literalai.my_types import PaginatedResponse 5 | from literalai.observability.filter import threads_filters, threads_order_by 6 | from literalai.observability.step import StepType 7 | from literalai.observability.thread import Thread 8 | 9 | 10 | def get_threads_helper( 11 | first: Optional[int] = None, 12 | after: Optional[str] = None, 13 | before: Optional[str] = None, 14 | filters: Optional[threads_filters] = None, 15 | order_by: Optional[threads_order_by] = None, 16 | step_types_to_keep: Optional[List[StepType]] = None, 17 | ): 18 | variables: Dict[str, Any] = {} 19 | 20 | if first: 21 | variables["first"] = first 22 | if after: 23 | variables["after"] = after 24 | if before: 25 | variables["before"] = before 26 | if filters: 27 | variables["filters"] = filters 28 | if order_by: 29 | variables["orderBy"] = order_by 30 | if step_types_to_keep: 31 | variables["stepTypesToKeep"] = step_types_to_keep 32 | 33 | def process_response(response): 34 | processed_response = response["data"]["threads"] 35 | processed_response["data"] = [ 36 | edge["node"] for edge in processed_response["edges"] 37 | ] 38 | return PaginatedResponse[Thread].from_dict(processed_response, Thread) 39 | 40 | description = "get threads" 41 | 42 | return gql.GET_THREADS, description, variables, process_response 43 | 44 | 45 | def list_threads_helper( 46 | first: Optional[int] = None, 47 | after: Optional[str] = None, 48 | before: Optional[str] = None, 49 | filters: Optional[threads_filters] = None, 50 | order_by: Optional[threads_order_by] = None, 51 | ): 52 | variables: Dict[str, Any] = {} 53 | 54 | if first: 55 | variables["first"] = first 56 | if after: 57 | variables["after"] = after 58 | if before: 59 | variables["before"] = before 60 | if filters: 61 | variables["filters"] = filters 62 | if order_by: 63 | variables["orderBy"] = order_by 64 | 65 | def process_response(response): 66 | response_data = response["data"]["threads"] 67 | response_data["data"] = list(map(lambda x: x["node"], response_data["edges"])) 68 | return PaginatedResponse[Thread].from_dict(response_data, Thread) 69 | 70 | description = "get threads" 71 | 72 | return gql.LIST_THREADS, description, variables, process_response 73 | 74 | 75 | def get_thread_helper(id: str): 76 | variables = {"id": id} 77 | 78 | def process_response(response): 79 | thread = response["data"]["threadDetail"] 80 | return Thread.from_dict(thread) if thread else None 81 | 82 | description = "get thread" 83 | 84 | return gql.GET_THREAD, description, variables, process_response 85 | 86 | 87 | def create_thread_helper( 88 | name: Optional[str] = None, 89 | metadata: Optional[Dict] = None, 90 | participant_id: Optional[str] = None, 91 | tags: Optional[List[str]] = None, 92 | ): 93 | variables = { 94 | "name": name, 95 | "metadata": metadata, 96 | "participantId": participant_id, 97 | "tags": tags, 98 | } 99 | 100 | def process_response(response): 101 | return Thread.from_dict(response["data"]["createThread"]) 102 | 103 | description = "create thread" 104 | 105 | return gql.CREATE_THREAD, description, variables, process_response 106 | 107 | 108 | def upsert_thread_helper( 109 | id: str, 110 | name: Optional[str] = None, 111 | metadata: Optional[Dict] = None, 112 | participant_id: Optional[str] = None, 113 | tags: Optional[List[str]] = None, 114 | ): 115 | variables = { 116 | "id": id, 117 | "name": name, 118 | "metadata": metadata, 119 | "participantId": participant_id, 120 | "tags": tags, 121 | } 122 | 123 | # remove None values to prevent the API from removing existing values 124 | variables = {k: v for k, v in variables.items() if v is not None} 125 | 126 | def process_response(response): 127 | return Thread.from_dict(response["data"]["upsertThread"]) 128 | 129 | description = "upsert thread" 130 | 131 | return gql.UPSERT_THREAD, description, variables, process_response 132 | 133 | 134 | def update_thread_helper( 135 | id: str, 136 | name: Optional[str] = None, 137 | metadata: Optional[Dict] = None, 138 | participant_id: Optional[str] = None, 139 | tags: Optional[List[str]] = None, 140 | ): 141 | variables = { 142 | "id": id, 143 | "name": name, 144 | "metadata": metadata, 145 | "participantId": participant_id, 146 | "tags": tags, 147 | } 148 | 149 | # remove None values to prevent the API from removing existing values 150 | variables = {k: v for k, v in variables.items() if v is not None} 151 | 152 | def process_response(response): 153 | return Thread.from_dict(response["data"]["updateThread"]) 154 | 155 | description = "update thread" 156 | 157 | return gql.UPDATE_THREAD, description, variables, process_response 158 | 159 | 160 | def delete_thread_helper(id: str): 161 | variables = {"thread_id": id} 162 | 163 | def process_response(response): 164 | deleted = bool(response["data"]["deleteThread"]) 165 | return deleted 166 | 167 | description = "delete thread" 168 | 169 | # Assuming DELETE_THREAD is a placeholder for the actual GraphQL mutation 170 | return gql.DELETE_THREAD, description, variables, process_response 171 | -------------------------------------------------------------------------------- /literalai/api/helpers/user_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | from literalai.api.helpers import gql 4 | from literalai.my_types import PaginatedResponse, User 5 | from literalai.observability.filter import users_filters 6 | 7 | 8 | def get_users_helper( 9 | first: Optional[int] = None, 10 | after: Optional[str] = None, 11 | before: Optional[str] = None, 12 | filters: Optional[users_filters] = None, 13 | ): 14 | variables: Dict[str, Any] = {} 15 | 16 | if first: 17 | variables["first"] = first 18 | if after: 19 | variables["after"] = after 20 | if before: 21 | variables["before"] = before 22 | if filters: 23 | variables["filters"] = filters 24 | 25 | def process_response(response): 26 | response = response["data"]["participants"] 27 | response["data"] = list(map(lambda x: x["node"], response["edges"])) 28 | return PaginatedResponse[User].from_dict(response, User) 29 | 30 | description = "get users" 31 | 32 | return gql.GET_PARTICIPANTS, description, variables, process_response 33 | 34 | 35 | def create_user_helper(identifier: str, metadata: Optional[Dict] = None): 36 | variables = {"identifier": identifier, "metadata": metadata} 37 | 38 | def process_response(response) -> User: 39 | return User.from_dict(response["data"]["createParticipant"]) 40 | 41 | description = "create user" 42 | 43 | return gql.CREATE_PARTICIPANT, description, variables, process_response 44 | 45 | 46 | def update_user_helper( 47 | id: str, identifier: Optional[str] = None, metadata: Optional[Dict] = None 48 | ): 49 | variables = {"id": id, "identifier": identifier, "metadata": metadata} 50 | 51 | # remove None values to prevent the API from removing existing values 52 | variables = {k: v for k, v in variables.items() if v is not None} 53 | 54 | def process_response(response): 55 | return User.from_dict(response["data"]["updateParticipant"]) 56 | 57 | description = "update user" 58 | 59 | return gql.UPDATE_PARTICIPANT, description, variables, process_response 60 | 61 | 62 | def get_user_helper(id: Optional[str] = None, identifier: Optional[str] = None): 63 | if id is None and identifier is None: 64 | raise Exception("Either id or identifier must be provided") 65 | 66 | if id is not None and identifier is not None: 67 | raise Exception("Only one of id or identifier must be provided") 68 | 69 | variables = {"id": id, "identifier": identifier} 70 | 71 | def process_response(response): 72 | user = response["data"]["participant"] 73 | return User.from_dict(user) if user else None 74 | 75 | description = "get user" 76 | 77 | return gql.GET_PARTICIPANT, description, variables, process_response 78 | 79 | 80 | def delete_user_helper(id: str): 81 | variables = {"id": id} 82 | 83 | def process_response(response): 84 | return response["data"]["deleteParticipant"]["id"] 85 | 86 | description = "delete user" 87 | 88 | return gql.DELETE_PARTICIPANT, description, variables, process_response 89 | -------------------------------------------------------------------------------- /literalai/cache/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/cache/__init__.py -------------------------------------------------------------------------------- /literalai/cache/prompt_helpers.py: -------------------------------------------------------------------------------- 1 | from literalai.cache.shared_cache import SharedCache 2 | from literalai.prompt_engineering.prompt import Prompt 3 | 4 | 5 | def put_prompt(cache: SharedCache, prompt: Prompt): 6 | cache.put(prompt.id, prompt) 7 | cache.put(prompt.name, prompt) 8 | cache.put(f"{prompt.name}-{prompt.version}", prompt) 9 | -------------------------------------------------------------------------------- /literalai/cache/shared_cache.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Optional 2 | 3 | 4 | class SharedCache: 5 | """ 6 | Singleton cache for storing data. 7 | Only one instance will exist regardless of how many times it's instantiated. 8 | """ 9 | 10 | _instance = None 11 | _cache: dict[str, Any] 12 | 13 | def __new__(cls): 14 | if cls._instance is None: 15 | cls._instance = super().__new__(cls) 16 | cls._instance._cache = {} 17 | return cls._instance 18 | 19 | def get_cache(self) -> dict[str, Any]: 20 | return self._cache 21 | 22 | def get(self, key: str) -> Optional[Any]: 23 | """ 24 | Retrieves a value from the cache using the provided key. 25 | """ 26 | if not isinstance(key, str): 27 | raise TypeError("Key must be a string") 28 | return self._cache.get(key) 29 | 30 | def put(self, key: str, value: Any): 31 | """ 32 | Stores a value in the cache. 33 | """ 34 | if not isinstance(key, str): 35 | raise TypeError("Key must be a string") 36 | self._cache[key] = value 37 | 38 | def clear(self) -> None: 39 | """ 40 | Clears all cached values. 41 | """ 42 | self._cache.clear() 43 | -------------------------------------------------------------------------------- /literalai/callback/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/callback/__init__.py -------------------------------------------------------------------------------- /literalai/context.py: -------------------------------------------------------------------------------- 1 | from contextvars import ContextVar 2 | from typing import TYPE_CHECKING, List, Optional 3 | 4 | if TYPE_CHECKING: 5 | from literalai.observability.step import Step 6 | from literalai.observability.thread import Thread 7 | 8 | active_steps_var = ContextVar[Optional[List["Step"]]]("active_steps", default=None) 9 | active_thread_var = ContextVar[Optional["Thread"]]("active_thread", default=None) 10 | active_root_run_var = ContextVar[Optional["Step"]]("active_root_run_var", default=None) 11 | 12 | active_experiment_item_run_id_var = ContextVar[Optional[str]]( 13 | "active_experiment_item_run", default=None 14 | ) 15 | -------------------------------------------------------------------------------- /literalai/environment.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import os 3 | from functools import wraps 4 | from typing import TYPE_CHECKING, Callable, Optional 5 | 6 | from literalai.my_types import Environment 7 | 8 | if TYPE_CHECKING: 9 | from literalai.client import BaseLiteralClient 10 | 11 | 12 | class EnvContextManager: 13 | def __init__(self, client: "BaseLiteralClient", env: Environment = "prod"): 14 | self.client = client 15 | self.env = env 16 | self.original_env = os.environ.get("LITERAL_ENV", "") 17 | 18 | def __call__(self, func): 19 | return env_decorator( 20 | self.client, 21 | func=func, 22 | ctx_manager=self, 23 | ) 24 | 25 | async def __aenter__(self): 26 | os.environ["LITERAL_ENV"] = self.env 27 | 28 | async def __aexit__(self, exc_type, exc_val, exc_tb): 29 | os.environ = self.original_env 30 | 31 | def __enter__(self): 32 | os.environ["LITERAL_ENV"] = self.env 33 | 34 | def __exit__(self, exc_type, exc_val, exc_tb): 35 | os.environ["LITERAL_ENV"] = self.original_env 36 | 37 | 38 | def env_decorator( 39 | client: "BaseLiteralClient", 40 | func: Callable, 41 | env: Environment = "prod", 42 | ctx_manager: Optional[EnvContextManager] = None, 43 | **decorator_kwargs, 44 | ): 45 | if not ctx_manager: 46 | ctx_manager = EnvContextManager( 47 | client=client, 48 | env=env, 49 | **decorator_kwargs, 50 | ) 51 | 52 | # Handle async decorator 53 | if inspect.iscoroutinefunction(func): 54 | 55 | @wraps(func) 56 | async def async_wrapper(*args, **kwargs): 57 | with ctx_manager: 58 | result = await func(*args, **kwargs) 59 | return result 60 | 61 | return async_wrapper 62 | else: 63 | # Handle sync decorator 64 | @wraps(func) 65 | def sync_wrapper(*args, **kwargs): 66 | with ctx_manager: 67 | return func(*args, **kwargs) 68 | 69 | return sync_wrapper 70 | -------------------------------------------------------------------------------- /literalai/evaluation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/evaluation/__init__.py -------------------------------------------------------------------------------- /literalai/evaluation/dataset.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from typing import TYPE_CHECKING, Dict, List, Literal, Optional, cast 3 | 4 | from typing_extensions import TypedDict 5 | 6 | from literalai.my_types import Utils 7 | 8 | if TYPE_CHECKING: 9 | from literalai.api import LiteralAPI 10 | 11 | from literalai.evaluation.dataset_experiment import DatasetExperiment 12 | from literalai.evaluation.dataset_item import DatasetItem, DatasetItemDict 13 | 14 | DatasetType = Literal["key_value", "generation"] 15 | 16 | 17 | class DatasetDict(TypedDict, total=False): 18 | id: str 19 | createdAt: str 20 | metadata: Dict 21 | name: Optional[str] 22 | description: Optional[str] 23 | items: Optional[List[DatasetItemDict]] 24 | type: DatasetType 25 | 26 | 27 | @dataclass(repr=False) 28 | class Dataset(Utils): 29 | """ 30 | A dataset of items, each item representing an ideal scenario to run experiments on. 31 | """ 32 | 33 | api: "LiteralAPI" 34 | id: str 35 | created_at: str 36 | metadata: Dict 37 | name: Optional[str] = None 38 | description: Optional[str] = None 39 | items: List[DatasetItem] = field(default_factory=lambda: []) 40 | type: DatasetType = "key_value" 41 | 42 | def to_dict(self): 43 | return { 44 | "id": self.id, 45 | "createdAt": self.created_at, 46 | "metadata": self.metadata, 47 | "name": self.name, 48 | "description": self.description, 49 | "items": [item.to_dict() for item in self.items], 50 | "type": self.type, 51 | } 52 | 53 | @classmethod 54 | def from_dict(cls, api: "LiteralAPI", dataset: DatasetDict) -> "Dataset": 55 | items = dataset.get("items", []) 56 | if not isinstance(items, list): 57 | raise Exception("Dataset items should be an array") 58 | 59 | return cls( 60 | api=api, 61 | id=dataset.get("id", ""), 62 | created_at=dataset.get("createdAt", ""), 63 | metadata=dataset.get("metadata", {}), 64 | name=dataset.get("name"), 65 | description=dataset.get("description"), 66 | items=[DatasetItem.from_dict(item) for item in items], 67 | type=dataset.get("type", cast(DatasetType, "key_value")), 68 | ) 69 | 70 | def update( 71 | self, 72 | name: Optional[str] = None, 73 | description: Optional[str] = None, 74 | metadata: Optional[Dict] = None, 75 | ): 76 | """ 77 | Update the dataset with the given name, description and metadata. 78 | """ 79 | updated_dataset = self.api.update_dataset( 80 | self.id, name=name, description=description, metadata=metadata 81 | ) 82 | self.name = updated_dataset.name 83 | self.description = updated_dataset.description 84 | self.metadata = updated_dataset.metadata 85 | 86 | def delete(self): 87 | """ 88 | Deletes the dataset. 89 | """ 90 | self.api.delete_dataset(self.id) 91 | 92 | def create_item( 93 | self, 94 | input: Dict, 95 | expected_output: Optional[Dict] = None, 96 | metadata: Optional[Dict] = None, 97 | ) -> DatasetItem: 98 | """ 99 | Creates a new dataset item and adds it to this dataset. 100 | 101 | Args: 102 | input: The input data for the dataset item. 103 | expected_output: The output data for the dataset item (optional). 104 | metadata: Metadata for the dataset item (optional). 105 | 106 | Returns: 107 | `DatasetItem`:The created DatasetItem instance. 108 | """ 109 | dataset_item = self.api.create_dataset_item( 110 | self.id, input, expected_output, metadata 111 | ) 112 | if self.items is None: 113 | self.items = [] 114 | self.items.append(dataset_item) 115 | return dataset_item 116 | 117 | def create_experiment( 118 | self, 119 | name: str, 120 | prompt_variant_id: Optional[str] = None, 121 | params: Optional[Dict] = None, 122 | ) -> DatasetExperiment: 123 | """ 124 | Creates a new dataset experiment based on this dataset. 125 | 126 | Args: 127 | name: The name of the experiment. 128 | prompt_variant_id: The Prompt variant ID to experiment on. 129 | params: The params used on the experiment. 130 | 131 | Returns: 132 | `DatasetExperiment`: The created DatasetExperiment instance. 133 | """ 134 | experiment = self.api.create_experiment( 135 | name=name, 136 | dataset_id=self.id, 137 | prompt_variant_id=prompt_variant_id, 138 | params=params, 139 | ) 140 | return experiment 141 | 142 | def delete_item(self, item_id: str): 143 | """ 144 | Deletes a dataset item from this dataset. 145 | 146 | Args: 147 | item_id: The ID of the dataset item to delete. 148 | """ 149 | self.api.delete_dataset_item(item_id) 150 | if self.items is not None: 151 | self.items = [item for item in self.items if item.id != item_id] 152 | 153 | def add_step(self, step_id: str, metadata: Optional[Dict] = None) -> DatasetItem: 154 | """ 155 | Create a new dataset item based on a step and add it to this dataset. 156 | 157 | Args: 158 | step_id: The id of the step to add to the dataset. 159 | metadata: Metadata for the dataset item (optional). 160 | 161 | Returns: 162 | `DatasetItem`: The created DatasetItem instance. 163 | """ 164 | if self.type == "generation": 165 | raise ValueError("Cannot add a step to a generation dataset") 166 | 167 | dataset_item = self.api.add_step_to_dataset(self.id, step_id, metadata) 168 | if self.items is None: 169 | self.items = [] 170 | self.items.append(dataset_item) 171 | return dataset_item 172 | 173 | def add_generation( 174 | self, generation_id: str, metadata: Optional[Dict] = None 175 | ) -> DatasetItem: 176 | """ 177 | Create a new dataset item based on a generation and add it to this dataset. 178 | 179 | Args: 180 | generation_id: The id of the generation to add to the dataset. 181 | metadata: Metadata for the dataset item (optional). 182 | 183 | Returns: 184 | `DatasetItem`: The created DatasetItem instance. 185 | """ 186 | dataset_item = self.api.add_generation_to_dataset( 187 | self.id, generation_id, metadata 188 | ) 189 | if self.items is None: 190 | self.items = [] 191 | self.items.append(dataset_item) 192 | return dataset_item 193 | -------------------------------------------------------------------------------- /literalai/evaluation/dataset_experiment.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from typing import TYPE_CHECKING, Dict, List, Optional, TypedDict 3 | 4 | from literalai.context import active_experiment_item_run_id_var 5 | from literalai.my_types import Utils 6 | from literalai.observability.step import ScoreDict 7 | 8 | if TYPE_CHECKING: 9 | from literalai.api import LiteralAPI 10 | 11 | 12 | class DatasetExperimentItemDict(TypedDict, total=False): 13 | id: str 14 | datasetExperimentId: str 15 | datasetItemId: Optional[str] 16 | scores: List[ScoreDict] 17 | input: Optional[Dict] 18 | output: Optional[Dict] 19 | experimentRunId: Optional[str] 20 | 21 | 22 | @dataclass(repr=False) 23 | class DatasetExperimentItem(Utils): 24 | """ 25 | An item of a `DatasetExperiment`: it may be linked to a `DatasetItem`. 26 | """ 27 | 28 | id: str 29 | dataset_experiment_id: str 30 | dataset_item_id: Optional[str] 31 | scores: List[ScoreDict] 32 | input: Optional[Dict] 33 | output: Optional[Dict] 34 | experiment_run_id: Optional[str] 35 | 36 | def to_dict(self): 37 | return { 38 | "id": self.id, 39 | "datasetExperimentId": self.dataset_experiment_id, 40 | "datasetItemId": self.dataset_item_id, 41 | "experimentRunId": self.experiment_run_id, 42 | "scores": self.scores, 43 | "input": self.input, 44 | "output": self.output, 45 | } 46 | 47 | @classmethod 48 | def from_dict(cls, item: DatasetExperimentItemDict) -> "DatasetExperimentItem": 49 | return cls( 50 | id=item.get("id", ""), 51 | experiment_run_id=item.get("experimentRunId"), 52 | dataset_experiment_id=item.get("datasetExperimentId", ""), 53 | dataset_item_id=item.get("datasetItemId"), 54 | scores=item.get("scores", []), 55 | input=item.get("input"), 56 | output=item.get("output"), 57 | ) 58 | 59 | 60 | class DatasetExperimentDict(TypedDict, total=False): 61 | id: str 62 | createdAt: str 63 | name: str 64 | datasetId: str 65 | params: Dict 66 | promptExperimentId: Optional[str] 67 | items: Optional[List[DatasetExperimentItemDict]] 68 | 69 | 70 | @dataclass(repr=False) 71 | class DatasetExperiment(Utils): 72 | """ 73 | An experiment, linked or not to a `Dataset`. 74 | """ 75 | 76 | api: "LiteralAPI" 77 | id: str 78 | created_at: str 79 | name: str 80 | dataset_id: Optional[str] 81 | params: Optional[Dict] 82 | prompt_variant_id: Optional[str] = None 83 | items: List[DatasetExperimentItem] = field(default_factory=lambda: []) 84 | 85 | def log(self, item_dict: DatasetExperimentItemDict) -> DatasetExperimentItem: 86 | """ 87 | Logs an item to the dataset experiment. 88 | """ 89 | experiment_run_id = active_experiment_item_run_id_var.get() 90 | dataset_experiment_item = DatasetExperimentItem.from_dict( 91 | { 92 | "experimentRunId": experiment_run_id, 93 | "datasetExperimentId": self.id, 94 | "datasetItemId": item_dict.get("datasetItemId"), 95 | "input": item_dict.get("input", {}), 96 | "output": item_dict.get("output", {}), 97 | "scores": item_dict.get("scores", []), 98 | } 99 | ) 100 | 101 | item = self.api.create_experiment_item(dataset_experiment_item) 102 | self.items.append(item) 103 | return item 104 | 105 | def to_dict(self): 106 | return { 107 | "id": self.id, 108 | "createdAt": self.created_at, 109 | "name": self.name, 110 | "datasetId": self.dataset_id, 111 | "promptExperimentId": self.prompt_variant_id, 112 | "params": self.params, 113 | "items": [item.to_dict() for item in self.items], 114 | } 115 | 116 | @classmethod 117 | def from_dict( 118 | cls, api: "LiteralAPI", dataset_experiment: DatasetExperimentDict 119 | ) -> "DatasetExperiment": 120 | items = dataset_experiment.get("items", []) 121 | if not isinstance(items, list): 122 | raise Exception("Dataset items should be a list.") 123 | return cls( 124 | api=api, 125 | id=dataset_experiment.get("id", ""), 126 | created_at=dataset_experiment.get("createdAt", ""), 127 | name=dataset_experiment.get("name", ""), 128 | dataset_id=dataset_experiment.get("datasetId", ""), 129 | params=dataset_experiment.get("params"), 130 | prompt_variant_id=dataset_experiment.get("promptExperimentId"), 131 | items=[DatasetExperimentItem.from_dict(item) for item in items], 132 | ) 133 | -------------------------------------------------------------------------------- /literalai/evaluation/dataset_item.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Dict, List, Optional, TypedDict 3 | 4 | from literalai.my_types import Utils 5 | 6 | 7 | class DatasetItemDict(TypedDict, total=False): 8 | id: str 9 | createdAt: str 10 | datasetId: str 11 | metadata: Dict 12 | input: Dict 13 | expectedOutput: Optional[Dict] 14 | intermediarySteps: Optional[List[Dict]] 15 | 16 | 17 | @dataclass(repr=False) 18 | class DatasetItem(Utils): 19 | """ 20 | A `Dataset` item, containing `input`, `expectedOutput` and `metadata`. 21 | """ 22 | 23 | id: str 24 | created_at: str 25 | dataset_id: str 26 | metadata: Dict 27 | input: Dict 28 | expected_output: Optional[Dict] = None 29 | intermediary_steps: Optional[List[Dict]] = None 30 | 31 | def to_dict(self): 32 | """ 33 | Dumps the contents of the object into a dictionary. 34 | """ 35 | return { 36 | "id": self.id, 37 | "createdAt": self.created_at, 38 | "datasetId": self.dataset_id, 39 | "metadata": self.metadata, 40 | "input": self.input, 41 | "expectedOutput": self.expected_output, 42 | "intermediarySteps": self.intermediary_steps, 43 | } 44 | 45 | @classmethod 46 | def from_dict(cls, dataset_item: DatasetItemDict) -> "DatasetItem": 47 | """ 48 | Builds a `DatasetItem` object from a dictionary. 49 | """ 50 | return cls( 51 | id=dataset_item.get("id", ""), 52 | created_at=dataset_item.get("createdAt", ""), 53 | dataset_id=dataset_item.get("datasetId", ""), 54 | metadata=dataset_item.get("metadata", {}), 55 | input=dataset_item.get("input", {}), 56 | expected_output=dataset_item.get("expectedOutput"), 57 | intermediary_steps=dataset_item.get("intermediarySteps"), 58 | ) 59 | -------------------------------------------------------------------------------- /literalai/evaluation/experiment_item_run.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import uuid 3 | from functools import wraps 4 | from typing import TYPE_CHECKING, Callable, Optional 5 | 6 | from literalai.context import active_experiment_item_run_id_var 7 | from literalai.environment import EnvContextManager 8 | from literalai.observability.step import StepContextManager 9 | 10 | if TYPE_CHECKING: 11 | from literalai.client import BaseLiteralClient 12 | 13 | 14 | class ExperimentItemRunContextManager(EnvContextManager, StepContextManager): 15 | def __init__( 16 | self, 17 | client: "BaseLiteralClient", 18 | ): 19 | self.client = client 20 | EnvContextManager.__init__(self, client=client, env="experiment") 21 | 22 | def __call__(self, func): 23 | return experiment_item_run_decorator( 24 | self.client, 25 | func=func, 26 | ctx_manager=self, 27 | ) 28 | 29 | async def __aenter__(self): 30 | id = str(uuid.uuid4()) 31 | StepContextManager.__init__( 32 | self, client=self.client, name="Experiment Run", type="run", id=id 33 | ) 34 | active_experiment_item_run_id_var.set(id) 35 | await EnvContextManager.__aenter__(self) 36 | step = await StepContextManager.__aenter__(self) 37 | return step 38 | 39 | async def __aexit__(self, exc_type, exc_val, exc_tb): 40 | await StepContextManager.__aexit__(self, exc_type, exc_val, exc_tb) 41 | await self.client.event_processor.aflush() 42 | await EnvContextManager.__aexit__(self, exc_type, exc_val, exc_tb) 43 | active_experiment_item_run_id_var.set(None) 44 | 45 | def __enter__(self): 46 | id = str(uuid.uuid4()) 47 | StepContextManager.__init__( 48 | self, client=self.client, name="Experiment Run", type="run", id=id 49 | ) 50 | active_experiment_item_run_id_var.set(id) 51 | EnvContextManager.__enter__(self) 52 | step = StepContextManager.__enter__(self) 53 | return step 54 | 55 | def __exit__(self, exc_type, exc_val, exc_tb): 56 | StepContextManager.__exit__(self, exc_type, exc_val, exc_tb) 57 | self.client.event_processor.flush() 58 | EnvContextManager.__exit__(self, exc_type, exc_val, exc_tb) 59 | active_experiment_item_run_id_var.set(None) 60 | 61 | 62 | def experiment_item_run_decorator( 63 | client: "BaseLiteralClient", 64 | func: Callable, 65 | ctx_manager: Optional[ExperimentItemRunContextManager] = None, 66 | **decorator_kwargs, 67 | ): 68 | if not ctx_manager: 69 | ctx_manager = ExperimentItemRunContextManager( 70 | client=client, 71 | **decorator_kwargs, 72 | ) 73 | 74 | # Handle async decorator 75 | if inspect.iscoroutinefunction(func): 76 | 77 | @wraps(func) 78 | async def async_wrapper(*args, **kwargs): 79 | with ctx_manager: 80 | result = await func(*args, **kwargs) 81 | return result 82 | 83 | return async_wrapper 84 | else: 85 | # Handle sync decorator 86 | @wraps(func) 87 | def sync_wrapper(*args, **kwargs): 88 | with ctx_manager: 89 | return func(*args, **kwargs) 90 | 91 | return sync_wrapper 92 | -------------------------------------------------------------------------------- /literalai/event_processor.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | import queue 4 | import threading 5 | import time 6 | import traceback 7 | from typing import TYPE_CHECKING, Callable, List, Optional 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | if TYPE_CHECKING: 12 | from literalai.api import LiteralAPI 13 | from literalai.observability.step import StepDict 14 | 15 | DEFAULT_SLEEP_TIME = 0.2 16 | 17 | 18 | # to_thread is a backport of asyncio.to_thread from Python 3.9 19 | async def to_thread(func, /, *args, **kwargs): 20 | import contextvars 21 | import functools 22 | 23 | loop = asyncio.get_running_loop() 24 | ctx = contextvars.copy_context() 25 | func_call = functools.partial(ctx.run, func, *args, **kwargs) 26 | return await loop.run_in_executor(None, func_call) 27 | 28 | 29 | class EventProcessor: 30 | event_queue: queue.Queue 31 | batch: List["StepDict"] 32 | batch_timeout: float = 5.0 33 | 34 | def __init__( 35 | self, 36 | api: "LiteralAPI", 37 | batch_size: int = 1, 38 | disabled: bool = False, 39 | preprocess_steps_function: Optional[ 40 | Callable[[List["StepDict"]], List["StepDict"]] 41 | ] = None, 42 | ): 43 | self.stop_event = threading.Event() 44 | self.batch_size = batch_size 45 | self.api = api 46 | self.event_queue = queue.Queue() 47 | self.disabled = disabled 48 | self.processing_counter = 0 49 | self.counter_lock = threading.Lock() 50 | self.last_batch_time = time.time() 51 | self.preprocess_steps_function = preprocess_steps_function 52 | self.processing_thread = threading.Thread( 53 | target=self._process_events, daemon=True 54 | ) 55 | if not self.disabled: 56 | self.processing_thread.start() 57 | 58 | def add_event(self, event: "StepDict"): 59 | with self.counter_lock: 60 | self.processing_counter += 1 61 | self.event_queue.put(event) 62 | 63 | async def a_add_events(self, event: "StepDict"): 64 | with self.counter_lock: 65 | self.processing_counter += 1 66 | await to_thread(self.event_queue.put, event) 67 | 68 | def set_preprocess_steps_function( 69 | self, 70 | preprocess_steps_function: Optional[ 71 | Callable[[List["StepDict"]], List["StepDict"]] 72 | ], 73 | ): 74 | """ 75 | Set a function that will preprocess steps before sending them to the API. 76 | The function should take a list of StepDict objects and return a list of processed StepDict objects. 77 | This can be used for tasks like PII removal. 78 | 79 | Args: 80 | preprocess_steps_function (Callable[[List["StepDict"]], List["StepDict"]]): The preprocessing function 81 | """ 82 | self.preprocess_steps_function = preprocess_steps_function 83 | 84 | def _process_events(self): 85 | while True: 86 | batch = [] 87 | start_time = time.time() 88 | try: 89 | elapsed_time = time.time() - start_time 90 | # Try to fill the batch up to the batch_size 91 | while ( 92 | len(batch) < self.batch_size and elapsed_time < self.batch_timeout 93 | ): 94 | # Attempt to get events with a timeout 95 | event = self.event_queue.get(timeout=0.5) 96 | batch.append(event) 97 | except queue.Empty: 98 | # No more events at the moment, proceed with processing what's in the batch 99 | pass 100 | 101 | # Process the batch if any events are present 102 | if batch: 103 | self._process_batch(batch) 104 | 105 | # Stop if the stop_event is set and the queue is empty 106 | if self.stop_event.is_set() and self.event_queue.empty(): 107 | break 108 | 109 | def _try_process_batch(self, batch: List): 110 | try: 111 | # Apply preprocessing function if it exists 112 | if self.preprocess_steps_function is not None: 113 | try: 114 | processed_batch = self.preprocess_steps_function(batch) 115 | # Only use the processed batch if it's valid 116 | if processed_batch is not None and isinstance( 117 | processed_batch, list 118 | ): 119 | batch = processed_batch 120 | else: 121 | logger.warning( 122 | "Preprocess function returned invalid result, using original batch" 123 | ) 124 | except Exception as e: 125 | logger.error(f"Error in preprocess function: {str(e)}") 126 | logger.error(traceback.format_exc()) 127 | # Continue with the original batch 128 | 129 | return self.api.send_steps(batch) 130 | except Exception: 131 | logger.error(f"Failed to send steps: {traceback.format_exc()}") 132 | return None 133 | 134 | def _process_batch(self, batch: List): 135 | # Simple one-try retry in case of network failure (no retry on graphql errors) 136 | retries = 0 137 | while not self._try_process_batch(batch) and retries < 1: 138 | retries += 1 139 | time.sleep(DEFAULT_SLEEP_TIME) 140 | with self.counter_lock: 141 | self.processing_counter -= len(batch) 142 | 143 | def flush_and_stop(self): 144 | self.stop_event.set() 145 | if not self.disabled: 146 | self.processing_thread.join() 147 | 148 | async def aflush(self): 149 | while not self.event_queue.empty() or self._is_processing(): 150 | await asyncio.sleep(DEFAULT_SLEEP_TIME) 151 | 152 | def flush(self): 153 | while not self.event_queue.empty() or self._is_processing(): 154 | time.sleep(DEFAULT_SLEEP_TIME) 155 | 156 | def _is_processing(self): 157 | with self.counter_lock: 158 | return self.processing_counter > 0 159 | 160 | def __del__(self): 161 | self.flush_and_stop() 162 | -------------------------------------------------------------------------------- /literalai/exporter.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | from datetime import datetime, timezone 4 | from typing import Dict, List, Optional, Sequence, Union, cast 5 | 6 | from opentelemetry.sdk.trace import ReadableSpan 7 | from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult 8 | 9 | from literalai.event_processor import EventProcessor 10 | from literalai.helper import utc_now 11 | from literalai.observability.generation import GenerationType 12 | from literalai.observability.step import Step, StepDict 13 | from literalai.prompt_engineering.prompt import PromptDict 14 | 15 | 16 | # TODO: Suppport Gemini models https://github.com/traceloop/openllmetry/issues/2419 17 | # TODO: Support llamaindex workflow https://github.com/traceloop/openllmetry/pull/2421 18 | class LoggingSpanExporter(SpanExporter): 19 | def __init__( 20 | self, 21 | logger_name: str = "span_exporter", 22 | event_processor: Optional[EventProcessor] = None, 23 | ): 24 | self.logger = logging.getLogger(logger_name) 25 | self.logger.setLevel(logging.INFO) 26 | self.event_processor = event_processor 27 | 28 | if not self.logger.handlers: 29 | handler = logging.StreamHandler() 30 | formatter = logging.Formatter( 31 | "%(asctime)s - %(name)s - %(levelname)s - %(message)s" 32 | ) 33 | handler.setFormatter(formatter) 34 | self.logger.addHandler(handler) 35 | 36 | def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: 37 | """Export the spans by logging them.""" 38 | try: 39 | for span in spans: 40 | if ( 41 | span.attributes 42 | and span.attributes.get("gen_ai.request.model", None) is not None 43 | and self.event_processor is not None 44 | ): 45 | step = self._create_step_from_span(span) 46 | self.event_processor.add_event(cast(StepDict, step.to_dict())) 47 | 48 | return SpanExportResult.SUCCESS 49 | except Exception as e: 50 | self.logger.error(f"Failed to export spans: {e}") 51 | return SpanExportResult.FAILURE 52 | 53 | def shutdown(self): 54 | """Shuts down the exporter.""" 55 | if self.event_processor is not None: 56 | return self.event_processor.flush_and_stop() 57 | 58 | def force_flush(self, timeout_millis: float = 30000) -> bool: 59 | """Force flush the exporter.""" 60 | return True 61 | 62 | def _create_step_from_span(self, span: ReadableSpan) -> Step: 63 | """Convert a span to a Step object""" 64 | attributes = span.attributes or {} 65 | 66 | start_time = ( 67 | datetime.fromtimestamp(span.start_time / 1e9, tz=timezone.utc).isoformat() 68 | if span.start_time 69 | else utc_now() 70 | ) 71 | end_time = ( 72 | datetime.fromtimestamp(span.end_time / 1e9, tz=timezone.utc).isoformat() 73 | if span.end_time 74 | else utc_now() 75 | ) 76 | duration, token_throughput = self._calculate_duration_and_throughput( 77 | span.start_time, 78 | span.end_time, 79 | int(str(attributes.get("llm.usage.total_tokens", 0))), 80 | ) 81 | 82 | generation_type = attributes.get("llm.request.type") 83 | is_chat = generation_type == "chat" 84 | 85 | span_props = { 86 | "parent_id": attributes.get( 87 | "traceloop.association.properties.literal.parent_id" 88 | ), 89 | "thread_id": attributes.get( 90 | "traceloop.association.properties.literal.thread_id" 91 | ), 92 | "root_run_id": attributes.get( 93 | "traceloop.association.properties.literal.root_run_id" 94 | ), 95 | "metadata": attributes.get( 96 | "traceloop.association.properties.literal.metadata" 97 | ), 98 | "tags": attributes.get("traceloop.association.properties.literal.tags"), 99 | "name": attributes.get("traceloop.association.properties.literal.name"), 100 | } 101 | 102 | span_props = { 103 | k: str(v) for k, v in span_props.items() if v is not None and v != "None" 104 | } 105 | 106 | serialized_prompt = attributes.get( 107 | "traceloop.association.properties.literal.prompt" 108 | ) 109 | prompt = cast( 110 | Optional[PromptDict], 111 | ( 112 | self._extract_json(str(serialized_prompt)) 113 | if serialized_prompt and serialized_prompt != "None" 114 | else None 115 | ), 116 | ) 117 | messages = self._extract_messages(cast(Dict, attributes)) if is_chat else [] 118 | 119 | message_completions = ( 120 | self._extract_messages(cast(Dict, attributes), "gen_ai.completion.") 121 | if is_chat 122 | else [] 123 | ) 124 | 125 | message_completion = message_completions[-1] if message_completions else None 126 | previous_messages = ( 127 | messages + message_completions[:-1] if message_completions else messages 128 | ) 129 | 130 | generation_content = { 131 | "duration": duration, 132 | "messages": previous_messages, 133 | "message_completion": message_completion, 134 | "prompt": attributes.get("gen_ai.prompt.0.user"), 135 | "promptId": prompt.get("id") if prompt else None, 136 | "completion": attributes.get("gen_ai.completion.0.content"), 137 | "model": attributes.get("gen_ai.request.model"), 138 | "provider": attributes.get("gen_ai.system"), 139 | "tokenThroughputInSeconds": token_throughput, 140 | "variables": prompt.get("variables") if prompt else None, 141 | } 142 | generation_settings = { 143 | "max_tokens": attributes.get("gen_ai.request.max_tokens"), 144 | "stream": attributes.get("llm.is_streaming"), 145 | "token_count": attributes.get("llm.usage.total_tokens"), 146 | "input_token_count": attributes.get("gen_ai.usage.prompt_tokens"), 147 | "output_token_count": attributes.get("gen_ai.usage.completion_tokens"), 148 | "frequency_penalty": attributes.get("gen_ai.request.frequency_penalty"), 149 | "presence_penalty": attributes.get("gen_ai.request.presence_penalty"), 150 | "temperature": attributes.get("gen_ai.request.temperature"), 151 | "top_p": attributes.get("gen_ai.request.top_p"), 152 | } 153 | 154 | step_dict = { 155 | "id": str(span.context.span_id) if span.context else None, 156 | "name": span_props.get("name", span.name), 157 | "type": "llm", 158 | "metadata": self._extract_json(str(span_props.get("metadata", "{}"))), 159 | "startTime": start_time, 160 | "endTime": end_time, 161 | "threadId": span_props.get("thread_id"), 162 | "parentId": span_props.get("parent_id"), 163 | "rootRunId": span_props.get("root_run_id"), 164 | "tags": self._extract_json(str(span_props.get("tags", "[]"))), 165 | "input": { 166 | "content": ( 167 | generation_content["messages"] 168 | if is_chat 169 | else generation_content["prompt"] 170 | ) 171 | }, 172 | "output": { 173 | "content": ( 174 | generation_content["message_completion"] 175 | if is_chat 176 | else generation_content["completion"] 177 | ) 178 | }, 179 | "generation": { 180 | "type": GenerationType.CHAT if is_chat else GenerationType.COMPLETION, 181 | "prompt": generation_content["prompt"] if not is_chat else None, 182 | "completion": generation_content["completion"] if not is_chat else None, 183 | "model": generation_content["model"], 184 | "provider": generation_content["provider"], 185 | "settings": generation_settings, 186 | "tokenCount": generation_settings["token_count"], 187 | "inputTokenCount": generation_settings["input_token_count"], 188 | "outputTokenCount": generation_settings["output_token_count"], 189 | "messages": generation_content["messages"], 190 | "messageCompletion": generation_content["message_completion"], 191 | }, 192 | } 193 | 194 | step = Step.from_dict(cast(StepDict, step_dict)) 195 | 196 | if not span.status.is_ok: 197 | step.error = span.status.description or "Unknown error" 198 | 199 | return step 200 | 201 | def _extract_messages( 202 | self, data: Dict, prefix: str = "gen_ai.prompt." 203 | ) -> List[Dict]: 204 | messages = [] 205 | index = 0 206 | 207 | while True: 208 | role_key = f"{prefix}{index}.role" 209 | content_key = f"{prefix}{index}.content" 210 | 211 | if role_key not in data or content_key not in data: 212 | break 213 | if data[role_key] == "placeholder": 214 | break 215 | 216 | messages.append( 217 | { 218 | "role": data[role_key], 219 | "content": self._extract_json(data[content_key]), 220 | } 221 | ) 222 | 223 | index += 1 224 | 225 | return messages 226 | 227 | def _extract_json(self, data: str) -> Union[Dict, List, str]: 228 | try: 229 | content = json.loads(data) 230 | except Exception: 231 | content = data 232 | 233 | return content 234 | 235 | def _calculate_duration_and_throughput( 236 | self, 237 | start_time_ns: Optional[int], 238 | end_time_ns: Optional[int], 239 | total_tokens: Optional[int], 240 | ) -> tuple[float, Optional[float]]: 241 | """Calculate duration in seconds and token throughput per second.""" 242 | duration_ns = ( 243 | end_time_ns - start_time_ns if start_time_ns and end_time_ns else 0 244 | ) 245 | duration_seconds = duration_ns / 1e9 246 | 247 | token_throughput = None 248 | if total_tokens is not None and duration_seconds > 0: 249 | token_throughput = total_tokens / duration_seconds 250 | 251 | return duration_seconds, token_throughput 252 | -------------------------------------------------------------------------------- /literalai/helper.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | from pydantic import BaseModel 4 | 5 | 6 | def filter_none_values(data): 7 | return {key: value for key, value in data.items() if value is not None} 8 | 9 | 10 | def ensure_values_serializable(data): 11 | """ 12 | Recursively ensures that all values in the input (dict or list) are JSON serializable. 13 | """ 14 | 15 | if isinstance(data, BaseModel): 16 | return filter_none_values(data.model_dump()) 17 | elif isinstance(data, dict): 18 | return {key: ensure_values_serializable(value) for key, value in data.items()} 19 | elif isinstance(data, list): 20 | return [ensure_values_serializable(item) for item in data] 21 | elif isinstance(data, (str, int, float, bool, type(None))): 22 | return data 23 | elif isinstance(data, (tuple, set)): 24 | return ensure_values_serializable( 25 | list(data) 26 | ) # Convert tuples and sets to lists 27 | else: 28 | return str(data) # Fallback: convert other types to string 29 | 30 | 31 | def force_dict(data, default_key="content"): 32 | if not isinstance(data, dict): 33 | return {default_key: data} 34 | return data 35 | 36 | 37 | def utc_now(): 38 | dt = datetime.utcnow() 39 | return dt.isoformat() + "Z" 40 | 41 | 42 | def timestamp_utc(timestamp: float): 43 | dt = datetime.utcfromtimestamp(timestamp) 44 | return dt.isoformat() + "Z" 45 | -------------------------------------------------------------------------------- /literalai/instrumentation/__init__.py: -------------------------------------------------------------------------------- 1 | MISTRALAI_PROVIDER = "mistralai" 2 | OPENAI_PROVIDER = "openai" 3 | -------------------------------------------------------------------------------- /literalai/instrumentation/llamaindex/__init__.py: -------------------------------------------------------------------------------- 1 | from llama_index.core.instrumentation import get_dispatcher 2 | 3 | from literalai.client import LiteralClient 4 | from literalai.instrumentation.llamaindex.event_handler import LiteralEventHandler 5 | from literalai.instrumentation.llamaindex.span_handler import LiteralSpanHandler 6 | 7 | is_llamaindex_instrumented = False 8 | 9 | 10 | def instrument_llamaindex(client: "LiteralClient"): 11 | """ 12 | Instruments LlamaIndex to automatically send logs to Literal AI. 13 | """ 14 | global is_llamaindex_instrumented 15 | if is_llamaindex_instrumented: 16 | return 17 | 18 | root_dispatcher = get_dispatcher() 19 | 20 | span_handler = LiteralSpanHandler() 21 | root_dispatcher.add_span_handler(span_handler) 22 | 23 | event_handler = LiteralEventHandler( 24 | literal_client=client, llama_index_span_handler=span_handler 25 | ) 26 | root_dispatcher.add_event_handler(event_handler) 27 | 28 | is_llamaindex_instrumented = True 29 | -------------------------------------------------------------------------------- /literalai/instrumentation/llamaindex/span_handler.py: -------------------------------------------------------------------------------- 1 | import uuid 2 | from typing import Any, Dict, Optional 3 | 4 | from llama_index.core.instrumentation.span import SimpleSpan 5 | from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler 6 | from llama_index.core.query_engine import RetrieverQueryEngine 7 | from typing_extensions import TypedDict 8 | 9 | from literalai.context import active_thread_var 10 | 11 | literalai_uuid_namespace = uuid.UUID("05f6b2b5-a912-47bd-958f-98a9c4496322") 12 | 13 | 14 | class SpanEntry(TypedDict): 15 | id: str 16 | parent_id: Optional[str] 17 | root_id: Optional[str] 18 | is_run_root: bool 19 | 20 | 21 | class LiteralSpanHandler(BaseSpanHandler[SimpleSpan]): 22 | """This class handles spans coming from LlamaIndex.""" 23 | 24 | spans: Dict[str, SpanEntry] = {} 25 | 26 | def __init__(self): 27 | super().__init__() 28 | 29 | def new_span( 30 | self, 31 | id_: str, 32 | bound_args: Any, 33 | instance: Optional[Any] = None, 34 | parent_span_id: Optional[str] = None, 35 | tags: Optional[Dict[str, Any]] = None, 36 | **kwargs: Any, 37 | ): 38 | self.spans[id_] = { 39 | "id": id_, 40 | "parent_id": parent_span_id, 41 | "root_id": None, 42 | "is_run_root": self.is_run_root(instance, parent_span_id), 43 | } 44 | 45 | if parent_span_id is not None: 46 | self.spans[id_]["root_id"] = self.get_root_span_id(parent_span_id) 47 | else: 48 | self.spans[id_]["root_id"] = id_ 49 | 50 | def prepare_to_exit_span( 51 | self, 52 | id_: str, 53 | bound_args: Any, 54 | instance: Optional[Any] = None, 55 | result: Optional[Any] = None, 56 | **kwargs: Any, 57 | ): 58 | """Logic for preparing to exit a span.""" 59 | if id_ in self.spans: 60 | del self.spans[id_] 61 | 62 | def prepare_to_drop_span( 63 | self, 64 | id_: str, 65 | bound_args: Any, 66 | instance: Optional[Any] = None, 67 | err: Optional[BaseException] = None, 68 | **kwargs: Any, 69 | ): 70 | """Logic for preparing to drop a span.""" 71 | if id_ in self.spans: 72 | del self.spans[id_] 73 | 74 | def is_run_root( 75 | self, instance: Optional[Any], parent_span_id: Optional[str] 76 | ) -> bool: 77 | """Returns True if the span is of type RetrieverQueryEngine, and it has no run root in its parent chain""" 78 | if not isinstance(instance, RetrieverQueryEngine): 79 | return False 80 | 81 | # Span is of correct type, we check that it doesn't have a run root in its parent chain 82 | while parent_span_id: 83 | parent_span = self.spans.get(parent_span_id) 84 | 85 | if not parent_span: 86 | parent_span_id = None 87 | continue 88 | 89 | if parent_span["is_run_root"]: 90 | return False 91 | 92 | parent_span_id = parent_span["parent_id"] 93 | 94 | return True 95 | 96 | def get_root_span_id(self, span_id: Optional[str]): 97 | """Finds the root span and returns its ID""" 98 | if not span_id: 99 | return None 100 | 101 | current_span = self.spans.get(span_id) 102 | 103 | if current_span is None: 104 | return None 105 | 106 | while current_span["parent_id"] is not None: 107 | current_span = self.spans.get(current_span["parent_id"]) 108 | if current_span is None: 109 | return None 110 | 111 | return current_span["id"] 112 | 113 | def get_run_id(self, span_id: Optional[str]): 114 | """Go up the span chain to find a run_root, return its ID (or None)""" 115 | if not span_id: 116 | return None 117 | 118 | current_span = self.spans.get(span_id) 119 | 120 | if current_span is None: 121 | return None 122 | 123 | while current_span: 124 | if current_span["is_run_root"]: 125 | return str(uuid.uuid5(literalai_uuid_namespace, current_span["id"])) 126 | 127 | parent_id = current_span["parent_id"] 128 | 129 | if parent_id: 130 | current_span = self.spans.get(parent_id) 131 | else: 132 | current_span = None 133 | 134 | return None 135 | 136 | def get_thread_id(self, span_id: Optional[str]): 137 | """Returns the root span ID as a thread ID""" 138 | active_thread = active_thread_var.get() 139 | 140 | if active_thread: 141 | return active_thread.id 142 | 143 | if span_id is None: 144 | return None 145 | 146 | current_span = self.spans.get(span_id) 147 | 148 | if current_span is None: 149 | return None 150 | 151 | root_id = current_span["root_id"] 152 | 153 | if not root_id: 154 | return None 155 | 156 | root_span = self.spans.get(root_id) 157 | 158 | if root_span is None: 159 | # span is already the root, uuid its own id 160 | return str(uuid.uuid5(literalai_uuid_namespace, span_id)) 161 | else: 162 | # uuid the id of the root span 163 | return str(uuid.uuid5(literalai_uuid_namespace, root_span["id"])) 164 | 165 | @classmethod 166 | def class_name(cls) -> str: 167 | """Class name.""" 168 | return "LiteralSpanHandler" 169 | -------------------------------------------------------------------------------- /literalai/my_types.py: -------------------------------------------------------------------------------- 1 | import json 2 | import uuid 3 | from abc import abstractmethod 4 | from typing import Any, Dict, Generic, List, Literal, Optional, Protocol, TypeVar 5 | 6 | from pydantic.dataclasses import Field, dataclass 7 | from typing_extensions import TypedDict 8 | 9 | Environment = Literal["dev", "staging", "prod", "experiment"] 10 | 11 | 12 | class Utils: 13 | def __str__(self): 14 | return json.dumps(self.to_dict(), sort_keys=True, indent=4) 15 | 16 | def __repr__(self): 17 | return json.dumps(self.to_dict(), sort_keys=True, indent=4) 18 | 19 | @abstractmethod 20 | def to_dict(self): 21 | pass 22 | 23 | 24 | @dataclass(repr=False) 25 | class PageInfo(Utils): 26 | has_next_page: bool 27 | start_cursor: Optional[str] 28 | end_cursor: Optional[str] 29 | 30 | def to_dict(self): 31 | return { 32 | "hasNextPage": self.has_next_page, 33 | "startCursor": self.start_cursor, 34 | "endCursor": self.end_cursor, 35 | } 36 | 37 | @classmethod 38 | def from_dict(cls, page_info_dict: Dict) -> "PageInfo": 39 | has_next_page = page_info_dict.get("hasNextPage", False) 40 | start_cursor = page_info_dict.get("startCursor", None) 41 | end_cursor = page_info_dict.get("endCursor", None) 42 | return cls( 43 | has_next_page=has_next_page, 44 | start_cursor=start_cursor, 45 | end_cursor=end_cursor, 46 | ) 47 | 48 | 49 | T = TypeVar("T", covariant=True) 50 | 51 | 52 | class HasFromDict(Protocol[T]): 53 | @classmethod 54 | def from_dict(cls, obj_dict: Any) -> T: 55 | raise NotImplementedError() 56 | 57 | 58 | @dataclass(repr=False) 59 | class PaginatedResponse(Generic[T], Utils): 60 | page_info: PageInfo 61 | data: List[T] 62 | total_count: Optional[int] = None 63 | 64 | def to_dict(self): 65 | return { 66 | "pageInfo": self.page_info.to_dict(), 67 | "totalCount": self.total_count, 68 | "data": [ 69 | (d.to_dict() if hasattr(d, "to_dict") and callable(d.to_dict) else d) 70 | for d in self.data 71 | ], 72 | } 73 | 74 | @classmethod 75 | def from_dict( 76 | cls, paginated_response_dict: Dict, the_class: HasFromDict[T] 77 | ) -> "PaginatedResponse[T]": 78 | page_info = PageInfo.from_dict(paginated_response_dict.get("pageInfo", {})) 79 | data = [the_class.from_dict(d) for d in paginated_response_dict.get("data", [])] 80 | total_count = paginated_response_dict.get("totalCount", None) 81 | return cls(page_info=page_info, data=data, total_count=total_count) 82 | 83 | 84 | class TextContent(TypedDict, total=False): 85 | type: Literal["text"] 86 | text: str 87 | 88 | 89 | class ImageUrlContent(TypedDict, total=False): 90 | type: Literal["image_url"] 91 | image_url: Dict 92 | 93 | 94 | class UserDict(TypedDict, total=False): 95 | id: Optional[str] 96 | metadata: Optional[Dict] 97 | identifier: Optional[str] 98 | createdAt: Optional[str] 99 | 100 | 101 | @dataclass(repr=False) 102 | class User(Utils): 103 | id: Optional[str] = None 104 | created_at: Optional[str] = None 105 | identifier: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4())) 106 | metadata: Dict = Field(default_factory=lambda: {}) 107 | 108 | def to_dict(self) -> UserDict: 109 | return { 110 | "id": self.id, 111 | "identifier": self.identifier, 112 | "metadata": self.metadata, 113 | "createdAt": self.created_at, 114 | } 115 | 116 | @classmethod 117 | def from_dict(cls, user_dict: Dict) -> "User": 118 | id_ = user_dict.get("id", "") 119 | identifier = user_dict.get("identifier", "") 120 | metadata = user_dict.get("metadata", {}) 121 | created_at = user_dict.get("createdAt", "") 122 | 123 | user = cls( 124 | id=id_, identifier=identifier, metadata=metadata, created_at=created_at 125 | ) 126 | 127 | return user 128 | -------------------------------------------------------------------------------- /literalai/observability/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/observability/__init__.py -------------------------------------------------------------------------------- /literalai/observability/filter.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Generic, List, Literal, Optional, TypeVar, Union 2 | 3 | from typing_extensions import TypedDict 4 | 5 | Field = TypeVar("Field") 6 | Operators = TypeVar("Operators") 7 | Value = TypeVar("Value") 8 | 9 | BOOLEAN_OPERATORS = Literal["is", "nis"] 10 | STRING_OPERATORS = Literal["eq", "neq", "ilike", "nilike"] 11 | NUMBER_OPERATORS = Literal["eq", "neq", "gt", "gte", "lt", "lte"] 12 | STRING_LIST_OPERATORS = Literal["in", "nin"] 13 | DATETIME_OPERATORS = Literal["gte", "lte", "gt", "lt"] 14 | 15 | OPERATORS = Union[ 16 | BOOLEAN_OPERATORS, 17 | STRING_OPERATORS, 18 | NUMBER_OPERATORS, 19 | STRING_LIST_OPERATORS, 20 | DATETIME_OPERATORS, 21 | ] 22 | 23 | 24 | class Filter(Generic[Field], TypedDict, total=False): 25 | field: Field 26 | operator: OPERATORS 27 | value: Any 28 | path: Optional[str] 29 | 30 | 31 | class OrderBy(Generic[Field], TypedDict): 32 | column: Field 33 | direction: Literal["ASC", "DESC"] 34 | 35 | 36 | threads_filterable_fields = Literal[ 37 | "id", 38 | "createdAt", 39 | "name", 40 | "stepType", 41 | "stepName", 42 | "stepOutput", 43 | "metadata", 44 | "tokenCount", 45 | "tags", 46 | "participantId", 47 | "participantIdentifiers", 48 | "scoreValue", 49 | "duration", 50 | ] 51 | threads_orderable_fields = Literal["createdAt", "tokenCount"] 52 | threads_filters = List[Filter[threads_filterable_fields]] 53 | threads_order_by = OrderBy[threads_orderable_fields] 54 | 55 | steps_filterable_fields = Literal[ 56 | "id", 57 | "name", 58 | "input", 59 | "output", 60 | "participantIdentifier", 61 | "startTime", 62 | "endTime", 63 | "metadata", 64 | "parentId", 65 | "threadId", 66 | "error", 67 | "tags", 68 | ] 69 | steps_orderable_fields = Literal["createdAt"] 70 | steps_filters = List[Filter[steps_filterable_fields]] 71 | steps_order_by = OrderBy[steps_orderable_fields] 72 | 73 | users_filterable_fields = Literal[ 74 | "id", 75 | "createdAt", 76 | "identifier", 77 | "lastEngaged", 78 | "threadCount", 79 | "tokenCount", 80 | "metadata", 81 | ] 82 | users_filters = List[Filter[users_filterable_fields]] 83 | 84 | scores_filterable_fields = Literal[ 85 | "id", 86 | "createdAt", 87 | "participant", 88 | "name", 89 | "tags", 90 | "value", 91 | "type", 92 | "comment", 93 | ] 94 | scores_orderable_fields = Literal["createdAt"] 95 | scores_filters = List[Filter[scores_filterable_fields]] 96 | scores_order_by = OrderBy[scores_orderable_fields] 97 | 98 | generation_filterable_fields = Literal[ 99 | "id", 100 | "createdAt", 101 | "model", 102 | "duration", 103 | "promptLineage", 104 | "promptVersion", 105 | "tags", 106 | "score", 107 | "participant", 108 | "tokenCount", 109 | "error", 110 | ] 111 | generation_orderable_fields = Literal[ 112 | "createdAt", 113 | "tokenCount", 114 | "model", 115 | "provider", 116 | "participant", 117 | "duration", 118 | ] 119 | generations_filters = List[Filter[generation_filterable_fields]] 120 | generations_order_by = OrderBy[generation_orderable_fields] 121 | -------------------------------------------------------------------------------- /literalai/observability/generation.py: -------------------------------------------------------------------------------- 1 | from enum import Enum, unique 2 | from typing import Dict, List, Literal, Optional, Union 3 | 4 | from pydantic import Field 5 | from pydantic.dataclasses import dataclass 6 | from typing_extensions import TypedDict 7 | 8 | from literalai.my_types import ImageUrlContent, TextContent, Utils 9 | 10 | GenerationMessageRole = Literal["user", "assistant", "tool", "function", "system"] 11 | 12 | 13 | @unique 14 | class GenerationType(str, Enum): 15 | CHAT = "CHAT" 16 | COMPLETION = "COMPLETION" 17 | 18 | def __str__(self): 19 | return self.value 20 | 21 | def __repr__(self): 22 | return f"GenerationType.{self.name}" 23 | 24 | def to_json(self): 25 | return self.value 26 | 27 | 28 | class GenerationMessage(TypedDict, total=False): 29 | uuid: Optional[str] 30 | templated: Optional[bool] 31 | name: Optional[str] 32 | role: Optional[GenerationMessageRole] 33 | content: Optional[Union[str, List[Union[TextContent, ImageUrlContent]]]] 34 | function_call: Optional[Dict] 35 | tool_calls: Optional[List[Dict]] 36 | tool_call_id: Optional[str] 37 | 38 | 39 | @dataclass(repr=False) 40 | class BaseGeneration(Utils): 41 | """ 42 | Base class for generation objects, containing common attributes and methods. 43 | 44 | Attributes: 45 | id (Optional[str]): The unique identifier of the generation. 46 | prompt_id (Optional[str]): The unique identifier of the prompt associated with the generation. 47 | provider (Optional[str]): The provider of the generation. 48 | model (Optional[str]): The model used for the generation. 49 | error (Optional[str]): Any error message associated with the generation. 50 | settings (Optional[Dict]): Settings used for the generation. 51 | variables (Optional[Dict]): Variables used in the generation. 52 | tags (Optional[List[str]]): Tags associated with the generation. 53 | metadata (Optional[Dict]): Metadata associated with the generation. 54 | tools (Optional[List[Dict]]): Tools used in the generation. 55 | token_count (Optional[int]): The total number of tokens in the generation. 56 | input_token_count (Optional[int]): The number of input tokens in the generation. 57 | output_token_count (Optional[int]): The number of output tokens in the generation. 58 | tt_first_token (Optional[float]): Time to first token in the generation. 59 | token_throughput_in_s (Optional[float]): Token throughput in seconds. 60 | duration (Optional[float]): Duration of the generation. 61 | 62 | Methods: 63 | from_dict(cls, generation_dict: Dict) -> Union["ChatGeneration", "CompletionGeneration"]: 64 | Creates a generation object from a dictionary. 65 | to_dict(self) -> Dict: 66 | Converts the generation object to a dictionary. 67 | """ 68 | 69 | id: Optional[str] = None 70 | prompt_id: Optional[str] = None 71 | provider: Optional[str] = None 72 | model: Optional[str] = None 73 | error: Optional[str] = None 74 | settings: Optional[Dict] = Field(default_factory=lambda: {}) 75 | variables: Optional[Dict] = Field(default_factory=lambda: {}) 76 | tags: Optional[List[str]] = Field(default_factory=lambda: []) 77 | metadata: Optional[Dict] = Field(default_factory=lambda: {}) 78 | tools: Optional[List[Dict]] = None 79 | token_count: Optional[int] = None 80 | input_token_count: Optional[int] = None 81 | output_token_count: Optional[int] = None 82 | tt_first_token: Optional[float] = None 83 | token_throughput_in_s: Optional[float] = None 84 | duration: Optional[float] = None 85 | 86 | @classmethod 87 | def from_dict( 88 | cls, generation_dict: Dict 89 | ) -> Union["ChatGeneration", "CompletionGeneration"]: 90 | type = GenerationType(generation_dict.get("type")) 91 | if type == GenerationType.CHAT: 92 | return ChatGeneration.from_dict(generation_dict) 93 | elif type == GenerationType.COMPLETION: 94 | return CompletionGeneration.from_dict(generation_dict) 95 | else: 96 | raise ValueError(f"Unknown generation type: {type}") 97 | 98 | def to_dict(self): 99 | _dict = { 100 | "promptId": self.prompt_id, 101 | "provider": self.provider, 102 | "model": self.model, 103 | "error": self.error, 104 | "settings": self.settings, 105 | "variables": self.variables, 106 | "tags": self.tags, 107 | "metadata": self.metadata, 108 | "tools": self.tools, 109 | "tokenCount": self.token_count, 110 | "inputTokenCount": self.input_token_count, 111 | "outputTokenCount": self.output_token_count, 112 | "ttFirstToken": self.tt_first_token, 113 | "tokenThroughputInSeconds": self.token_throughput_in_s, 114 | "duration": self.duration, 115 | } 116 | if self.id: 117 | _dict["id"] = self.id 118 | return _dict 119 | 120 | 121 | @dataclass(repr=False) 122 | class CompletionGeneration(BaseGeneration, Utils): 123 | """ 124 | Represents a completion generation with a prompt and its corresponding completion. 125 | 126 | Attributes: 127 | prompt (Optional[str]): The prompt text for the generation. 128 | completion (Optional[str]): The generated completion text. 129 | type (GenerationType): The type of generation, which is set to GenerationType.COMPLETION. 130 | """ 131 | 132 | prompt: Optional[str] = None 133 | completion: Optional[str] = None 134 | type = GenerationType.COMPLETION 135 | 136 | def to_dict(self): 137 | _dict = super().to_dict() 138 | _dict.update( 139 | { 140 | "prompt": self.prompt, 141 | "completion": self.completion, 142 | "type": self.type.value, 143 | } 144 | ) 145 | return _dict 146 | 147 | @classmethod 148 | def from_dict(cls, generation_dict: Dict): 149 | return CompletionGeneration( 150 | id=generation_dict.get("id"), 151 | prompt_id=generation_dict.get("promptId"), 152 | error=generation_dict.get("error"), 153 | tags=generation_dict.get("tags"), 154 | provider=generation_dict.get("provider"), 155 | model=generation_dict.get("model"), 156 | variables=generation_dict.get("variables"), 157 | tools=generation_dict.get("tools"), 158 | settings=generation_dict.get("settings"), 159 | token_count=generation_dict.get("tokenCount"), 160 | input_token_count=generation_dict.get("inputTokenCount"), 161 | output_token_count=generation_dict.get("outputTokenCount"), 162 | tt_first_token=generation_dict.get("ttFirstToken"), 163 | token_throughput_in_s=generation_dict.get("tokenThroughputInSeconds"), 164 | duration=generation_dict.get("duration"), 165 | prompt=generation_dict.get("prompt"), 166 | completion=generation_dict.get("completion"), 167 | ) 168 | 169 | 170 | @dataclass(repr=False) 171 | class ChatGeneration(BaseGeneration, Utils): 172 | """ 173 | Represents a chat generation with a list of messages and a message completion. 174 | 175 | Attributes: 176 | messages (Optional[List[GenerationMessage]]): The list of messages in the chat generation. 177 | message_completion (Optional[GenerationMessage]): The completion message of the chat generation. 178 | type (GenerationType): The type of generation, which is set to GenerationType.CHAT. 179 | """ 180 | 181 | type = GenerationType.CHAT 182 | messages: Optional[List[GenerationMessage]] = Field(default_factory=lambda: []) 183 | message_completion: Optional[GenerationMessage] = None 184 | 185 | def to_dict(self): 186 | _dict = super().to_dict() 187 | _dict.update( 188 | { 189 | "messages": self.messages, 190 | "messageCompletion": self.message_completion, 191 | "type": self.type.value, 192 | } 193 | ) 194 | return _dict 195 | 196 | @classmethod 197 | def from_dict(self, generation_dict: Dict): 198 | return ChatGeneration( 199 | id=generation_dict.get("id"), 200 | prompt_id=generation_dict.get("promptId"), 201 | error=generation_dict.get("error"), 202 | tags=generation_dict.get("tags"), 203 | provider=generation_dict.get("provider"), 204 | model=generation_dict.get("model"), 205 | variables=generation_dict.get("variables"), 206 | tools=generation_dict.get("tools"), 207 | settings=generation_dict.get("settings"), 208 | token_count=generation_dict.get("tokenCount"), 209 | input_token_count=generation_dict.get("inputTokenCount"), 210 | output_token_count=generation_dict.get("outputTokenCount"), 211 | tt_first_token=generation_dict.get("ttFirstToken"), 212 | token_throughput_in_s=generation_dict.get("tokenThroughputInSeconds"), 213 | duration=generation_dict.get("duration"), 214 | messages=generation_dict.get("messages", []), 215 | message_completion=generation_dict.get("messageCompletion"), 216 | ) 217 | -------------------------------------------------------------------------------- /literalai/observability/message.py: -------------------------------------------------------------------------------- 1 | import uuid 2 | from typing import TYPE_CHECKING, Dict, List, Optional 3 | 4 | if TYPE_CHECKING: 5 | from literalai.event_processor import EventProcessor 6 | 7 | from literalai.context import active_root_run_var, active_steps_var, active_thread_var 8 | from literalai.helper import utc_now 9 | from literalai.my_types import Utils 10 | from literalai.observability.step import Attachment, MessageStepType, Score, StepDict 11 | 12 | 13 | class Message(Utils): 14 | id: Optional[str] = None 15 | name: Optional[str] = None 16 | type: Optional[MessageStepType] = None 17 | metadata: Optional[Dict] = {} 18 | parent_id: Optional[str] = None 19 | timestamp: Optional[str] = None 20 | content: str 21 | thread_id: Optional[str] = None 22 | root_run_id: Optional[str] = None 23 | tags: Optional[List[str]] = None 24 | created_at: Optional[str] = None 25 | 26 | scores: List[Score] = [] 27 | attachments: List[Attachment] = [] 28 | 29 | def __init__( 30 | self, 31 | content: str, 32 | id: Optional[str] = None, 33 | type: Optional[MessageStepType] = "assistant_message", 34 | name: Optional[str] = None, 35 | thread_id: Optional[str] = None, 36 | parent_id: Optional[str] = None, 37 | scores: List[Score] = [], 38 | attachments: List[Attachment] = [], 39 | metadata: Optional[Dict] = {}, 40 | timestamp: Optional[str] = None, 41 | tags: Optional[List[str]] = [], 42 | processor: Optional["EventProcessor"] = None, 43 | root_run_id: Optional[str] = None, 44 | ): 45 | from time import sleep 46 | 47 | sleep(0.001) 48 | 49 | self.id = id or str(uuid.uuid4()) 50 | if not timestamp: 51 | self.timestamp = utc_now() 52 | else: 53 | self.timestamp = timestamp 54 | self.name = name 55 | self.type = type 56 | self.content = content 57 | self.scores = scores 58 | self.attachments = attachments 59 | self.metadata = metadata 60 | self.tags = tags 61 | 62 | self.processor = processor 63 | 64 | # priority for thread_id: thread_id > parent_step.thread_id > active_thread 65 | self.thread_id = thread_id 66 | 67 | # priority for root_run_id: root_run_id > parent_step.root_run_id > active_root_run 68 | self.root_run_id = root_run_id 69 | 70 | # priority for parent_id: parent_id > parent_step.id 71 | self.parent_id = parent_id 72 | 73 | def end(self): 74 | active_steps = active_steps_var.get() 75 | 76 | if active_steps: 77 | parent_step = active_steps[-1] 78 | if not self.parent_id: 79 | self.parent_id = parent_step.id 80 | if not self.thread_id: 81 | self.thread_id = parent_step.thread_id 82 | if not self.root_run_id: 83 | self.root_run_id = parent_step.root_run_id 84 | 85 | if not self.thread_id: 86 | if active_thread := active_thread_var.get(): 87 | self.thread_id = active_thread.id 88 | 89 | if not self.root_run_id: 90 | if active_root_run := active_root_run_var.get(): 91 | self.root_run_id = active_root_run.id 92 | 93 | if not self.thread_id and not self.parent_id: 94 | raise Exception( 95 | "Message must be initialized with a thread_id or a parent id." 96 | ) 97 | 98 | if self.processor is None: 99 | raise Exception( 100 | "Message must be initialized with a processor to allow finalization." 101 | ) 102 | self.processor.add_event(self.to_dict()) 103 | 104 | def to_dict(self) -> "StepDict": 105 | # Create a correct step Dict from a message 106 | return { 107 | "id": self.id, 108 | "metadata": self.metadata, 109 | "parentId": self.parent_id, 110 | "startTime": self.timestamp, 111 | "endTime": self.timestamp, # startTime = endTime in Message 112 | "type": self.type, 113 | "threadId": self.thread_id, 114 | "output": { 115 | "content": self.content 116 | }, # no input, output = content in Message 117 | "name": self.name, 118 | "tags": self.tags, 119 | "scores": [score.to_dict() for score in self.scores], 120 | "attachments": [attachment.to_dict() for attachment in self.attachments], 121 | "rootRunId": self.root_run_id, 122 | } 123 | 124 | @classmethod 125 | def from_dict(cls, message_dict: Dict) -> "Message": 126 | id = message_dict.get("id", None) 127 | type = message_dict.get("type", None) 128 | thread_id = message_dict.get("threadId", None) 129 | root_run_id = message_dict.get("rootRunId", None) 130 | 131 | metadata = message_dict.get("metadata", None) 132 | parent_id = message_dict.get("parentId", None) 133 | timestamp = message_dict.get("startTime", None) 134 | content = message_dict.get("output", {}).get("content", "") 135 | name = message_dict.get("name", None) 136 | scores = message_dict.get("scores", []) 137 | attachments = message_dict.get("attachments", []) 138 | 139 | message = cls( 140 | id=id, 141 | metadata=metadata, 142 | parent_id=parent_id, 143 | timestamp=timestamp, 144 | type=type, 145 | thread_id=thread_id, 146 | content=content, 147 | name=name, 148 | scores=scores, 149 | attachments=attachments, 150 | root_run_id=root_run_id, 151 | ) 152 | 153 | message.created_at = message_dict.get("createdAt", None) 154 | 155 | return message 156 | -------------------------------------------------------------------------------- /literalai/observability/thread.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import logging 3 | import traceback 4 | import uuid 5 | from functools import wraps 6 | from typing import TYPE_CHECKING, Callable, Dict, List, Optional, TypedDict 7 | 8 | from traceloop.sdk import Traceloop 9 | 10 | from literalai.context import active_thread_var 11 | from literalai.my_types import UserDict, Utils 12 | from literalai.observability.step import Step, StepDict 13 | 14 | if TYPE_CHECKING: 15 | from literalai.client import BaseLiteralClient 16 | 17 | logger = logging.getLogger(__name__) 18 | 19 | 20 | class ThreadDict(TypedDict, total=False): 21 | id: Optional[str] 22 | name: Optional[str] 23 | metadata: Optional[Dict] 24 | tags: Optional[List[str]] 25 | createdAt: Optional[str] 26 | steps: Optional[List[StepDict]] 27 | participant: Optional[UserDict] 28 | 29 | 30 | class Thread(Utils): 31 | """ 32 | ## Using the `with` statement 33 | 34 | If you prefer to have more flexibility in logging Threads, you can use the `with` statement. You can create a thread and execute code within it using the `with` statement: 35 | 36 | 37 | ```python 38 | with literal_client.thread() as thread: 39 | # do something 40 | ``` 41 | 42 | 43 | You can also continue a thread by passing the thread id to the `thread` method: 44 | 45 | 46 | ```python 47 | 48 | previous_thread_id = "UUID" 49 | 50 | with literal_client.thread(thread_id=previous_thread_id) as thread: 51 | # do something 52 | ``` 53 | 54 | 55 | ## Using the Literal AI API client 56 | 57 | You can also create Threads using the `literal_client.api.create_thread()` method. 58 | 59 | 60 | ```python 61 | thread = literal_client.api.create_thread( 62 | participant_id="", 63 | environment="production", 64 | tags=["tag1", "tag2"], 65 | metadata={"key": "value"}, 66 | ) 67 | ``` 68 | 69 | 70 | ## Using Chainlit 71 | 72 | If you built your LLM application with Chainlit, you don't need to specify Threads in your code. Chainlit logs Threads for you by default. 73 | """ 74 | 75 | id: str 76 | name: Optional[str] 77 | metadata: Optional[Dict] 78 | tags: Optional[List[str]] 79 | steps: Optional[List[Step]] 80 | participant_id: Optional[str] 81 | participant_identifier: Optional[str] = None 82 | created_at: Optional[str] 83 | 84 | def __init__( 85 | self, 86 | id: str, 87 | steps: Optional[List[Step]] = [], 88 | name: Optional[str] = None, 89 | metadata: Optional[Dict] = {}, 90 | tags: Optional[List[str]] = [], 91 | participant_id: Optional[str] = None, 92 | ): 93 | self.id = id 94 | self.steps = steps 95 | self.name = name 96 | self.metadata = metadata 97 | self.tags = tags 98 | self.participant_id = participant_id 99 | 100 | def to_dict(self) -> ThreadDict: 101 | return { 102 | "id": self.id, 103 | "metadata": self.metadata, 104 | "tags": self.tags, 105 | "name": self.name, 106 | "steps": [step.to_dict() for step in self.steps] if self.steps else [], 107 | "participant": ( 108 | UserDict(id=self.participant_id, identifier=self.participant_identifier) 109 | if self.participant_id 110 | else UserDict() 111 | ), 112 | "createdAt": getattr(self, "created_at", None), 113 | } 114 | 115 | @classmethod 116 | def from_dict(cls, thread_dict: ThreadDict) -> "Thread": 117 | step_dict_list = thread_dict.get("steps", None) or [] 118 | id = thread_dict.get("id", None) or "" 119 | name = thread_dict.get("name", None) 120 | metadata = thread_dict.get("metadata", {}) 121 | tags = thread_dict.get("tags", []) 122 | steps = [Step.from_dict(step_dict) for step_dict in step_dict_list] 123 | participant = thread_dict.get("participant", None) 124 | participant_id = participant.get("id", None) if participant else None 125 | participant_identifier = ( 126 | participant.get("identifier", None) if participant else None 127 | ) 128 | created_at = thread_dict.get("createdAt", None) 129 | 130 | thread = cls( 131 | id=id, 132 | steps=steps, 133 | name=name, 134 | metadata=metadata, 135 | tags=tags, 136 | participant_id=participant_id, 137 | ) 138 | 139 | thread.created_at = created_at 140 | thread.participant_identifier = participant_identifier 141 | 142 | return thread 143 | 144 | 145 | class ThreadContextManager: 146 | def __init__( 147 | self, 148 | client: "BaseLiteralClient", 149 | thread_id: "Optional[str]" = None, 150 | name: "Optional[str]" = None, 151 | **kwargs, 152 | ): 153 | self.client = client 154 | self.thread_id = thread_id 155 | self.name = name 156 | self.kwargs = kwargs 157 | 158 | def upsert(self): 159 | if self.client.disabled: 160 | return 161 | 162 | thread = active_thread_var.get() 163 | thread_data = thread.to_dict() 164 | thread_data_to_upsert = { 165 | "id": thread_data["id"], 166 | "name": thread_data["name"], 167 | } 168 | 169 | metadata = { 170 | **(self.client.global_metadata or {}), 171 | **(thread_data.get("metadata") or {}), 172 | } 173 | if metadata: 174 | thread_data_to_upsert["metadata"] = metadata 175 | if tags := thread_data.get("tags"): 176 | thread_data_to_upsert["tags"] = tags 177 | if participant_id := thread_data.get("participant", {}).get("id"): 178 | thread_data_to_upsert["participant_id"] = participant_id 179 | 180 | try: 181 | self.client.to_sync().api.upsert_thread(**thread_data_to_upsert) 182 | except Exception: 183 | logger.error(f"Failed to upsert thread: {traceback.format_exc()}") 184 | 185 | def __call__(self, func): 186 | return thread_decorator( 187 | self.client, func=func, name=self.name, ctx_manager=self 188 | ) 189 | 190 | def __enter__(self) -> "Optional[Thread]": 191 | thread_id = self.thread_id if self.thread_id else str(uuid.uuid4()) 192 | active_thread_var.set(Thread(id=thread_id, name=self.name, **self.kwargs)) 193 | Traceloop.set_association_properties( 194 | { 195 | "literal.thread_id": thread_id, 196 | } 197 | ) 198 | return active_thread_var.get() 199 | 200 | def __exit__(self, exc_type, exc_val, exc_tb): 201 | if active_thread_var.get(): 202 | self.upsert() 203 | active_thread_var.set(None) 204 | 205 | async def __aenter__(self): 206 | thread_id = self.thread_id if self.thread_id else str(uuid.uuid4()) 207 | active_thread_var.set(Thread(id=thread_id, name=self.name, **self.kwargs)) 208 | Traceloop.set_association_properties( 209 | { 210 | "literal.thread_id": thread_id, 211 | } 212 | ) 213 | return active_thread_var.get() 214 | 215 | async def __aexit__(self, exc_type, exc_val, exc_tb): 216 | if active_thread_var.get(): 217 | self.upsert() 218 | active_thread_var.set(None) 219 | 220 | 221 | def thread_decorator( 222 | client: "BaseLiteralClient", 223 | func: Callable, 224 | thread_id: Optional[str] = None, 225 | name: Optional[str] = None, 226 | ctx_manager: Optional[ThreadContextManager] = None, 227 | **decorator_kwargs, 228 | ): 229 | if not ctx_manager: 230 | ctx_manager = ThreadContextManager( 231 | client, thread_id=thread_id, name=name, **decorator_kwargs 232 | ) 233 | if inspect.iscoroutinefunction(func): 234 | 235 | @wraps(func) 236 | async def async_wrapper(*args, **kwargs): 237 | with ctx_manager: 238 | result = await func(*args, **kwargs) 239 | return result 240 | 241 | return async_wrapper 242 | else: 243 | 244 | @wraps(func) 245 | def sync_wrapper(*args, **kwargs): 246 | with ctx_manager: 247 | return func(*args, **kwargs) 248 | 249 | return sync_wrapper 250 | -------------------------------------------------------------------------------- /literalai/prompt_engineering/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/prompt_engineering/__init__.py -------------------------------------------------------------------------------- /literalai/prompt_engineering/prompt.py: -------------------------------------------------------------------------------- 1 | import html 2 | from dataclasses import dataclass 3 | from importlib.metadata import version 4 | from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional 5 | 6 | import chevron 7 | from pydantic import Field 8 | from typing_extensions import TypedDict, deprecated 9 | 10 | if TYPE_CHECKING: 11 | from literalai.api import LiteralAPI 12 | 13 | from literalai.my_types import Utils 14 | from literalai.observability.generation import GenerationMessage, GenerationType 15 | 16 | 17 | class ProviderSettings(TypedDict, total=False): 18 | provider: str 19 | model: str 20 | frequency_penalty: float 21 | max_tokens: int 22 | presence_penalty: float 23 | stop: Optional[List[str]] 24 | temperature: float 25 | top_p: float 26 | 27 | 28 | class PromptVariable(TypedDict, total=False): 29 | name: str 30 | language: Literal["json", "plaintext"] 31 | 32 | 33 | class LiteralMessageDict(dict): 34 | def __init__(self, prompt_id: str, variables: Dict, *args, **kwargs): 35 | super().__init__(*args, **kwargs) # Initialize as a regular dict 36 | if "uuid" in self: 37 | uuid = self.pop("uuid") 38 | self.__literal_prompt__ = { 39 | "uuid": uuid, 40 | "prompt_id": prompt_id, 41 | "variables": variables, 42 | } 43 | 44 | 45 | class PromptDict(TypedDict, total=False): 46 | id: str 47 | lineage: Dict 48 | createdAt: str 49 | updatedAt: str 50 | type: "GenerationType" 51 | name: str 52 | version: int 53 | url: str 54 | versionDesc: Optional[str] 55 | templateMessages: List["GenerationMessage"] 56 | tools: Optional[List[Dict]] 57 | provider: str 58 | settings: ProviderSettings 59 | variables: List[PromptVariable] 60 | variablesDefaultValues: Optional[Dict[str, Any]] 61 | 62 | 63 | @dataclass(repr=False) 64 | class Prompt(Utils): 65 | """ 66 | Represents a version of a prompt template with variables, tools and settings. 67 | 68 | Attributes 69 | ---------- 70 | template_messages : List[GenerationMessage] 71 | The messages that make up the prompt. Messages can be of type `text` or `image`. 72 | Messages can reference variables. 73 | variables : List[PromptVariable] 74 | Variables exposed in the prompt. 75 | tools : Optional[List[Dict]] 76 | Tools LLM can pick from. 77 | settings : ProviderSettings 78 | LLM provider settings. 79 | 80 | Methods 81 | ------- 82 | format_messages(**kwargs: Any): 83 | Formats the prompt's template messages with the given variables. 84 | Variables may be passed as a dictionary or as keyword arguments. 85 | Keyword arguments take precedence over variables passed as a dictionary. 86 | """ 87 | 88 | api: "LiteralAPI" 89 | id: str 90 | created_at: str 91 | updated_at: str 92 | type: "GenerationType" 93 | name: str 94 | version: int 95 | url: str 96 | version_desc: Optional[str] 97 | template_messages: List["GenerationMessage"] 98 | tools: Optional[List[Dict]] 99 | provider: str 100 | settings: ProviderSettings 101 | variables: List[PromptVariable] 102 | variables_default_values: Optional[Dict[str, Any]] 103 | 104 | def to_dict(self) -> PromptDict: 105 | return { 106 | "id": self.id, 107 | "createdAt": self.created_at, 108 | "updatedAt": self.updated_at, 109 | "type": self.type, 110 | "name": self.name, 111 | "version": self.version, 112 | "url": self.url, 113 | "versionDesc": self.version_desc, 114 | "templateMessages": self.template_messages, # Assuming this is a list of dicts or similar serializable objects 115 | "tools": self.tools, 116 | "provider": self.provider, 117 | "settings": self.settings, 118 | "variables": self.variables, 119 | "variablesDefaultValues": self.variables_default_values, 120 | } 121 | 122 | @classmethod 123 | def from_dict(cls, api: "LiteralAPI", prompt_dict: PromptDict) -> "Prompt": 124 | # Create a Prompt instance from a dictionary (PromptDict) 125 | settings = prompt_dict.get("settings") or {} 126 | provider = settings.pop("provider", "") 127 | 128 | return cls( 129 | api=api, 130 | id=prompt_dict.get("id", ""), 131 | name=prompt_dict.get("lineage", {}).get("name", ""), 132 | version=prompt_dict.get("version", 0), 133 | url=prompt_dict.get("url", ""), 134 | created_at=prompt_dict.get("createdAt", ""), 135 | updated_at=prompt_dict.get("updatedAt", ""), 136 | type=prompt_dict.get("type", GenerationType.CHAT), 137 | version_desc=prompt_dict.get("versionDesc"), 138 | template_messages=prompt_dict.get("templateMessages", []), 139 | tools=prompt_dict.get("tools", []), 140 | provider=provider, 141 | settings=settings, 142 | variables=prompt_dict.get("variables", []), 143 | variables_default_values=prompt_dict.get("variablesDefaultValues"), 144 | ) 145 | 146 | def format_messages(self, **kwargs: Any) -> List[Any]: 147 | """ 148 | Formats the prompt's template messages with the given variables. 149 | Variables may be passed as a dictionary or as keyword arguments. 150 | Keyword arguments take precedence over variables passed as a dictionary. 151 | 152 | Args: 153 | variables (Optional[Dict[str, Any]]): Optional variables to resolve in the template messages. 154 | 155 | Returns: 156 | List[Any]: List of formatted chat completion messages. 157 | """ 158 | variables_with_defaults = { 159 | **(self.variables_default_values or {}), 160 | **(kwargs or {}), 161 | } 162 | formatted_messages = [] 163 | 164 | for message in self.template_messages: 165 | formatted_message = LiteralMessageDict( 166 | self.id, variables_with_defaults, message.copy() 167 | ) 168 | if isinstance(formatted_message["content"], str): 169 | formatted_message["content"] = html.unescape( 170 | chevron.render(message["content"], variables_with_defaults) 171 | ) 172 | else: 173 | for content in formatted_message["content"]: 174 | if content["type"] == "text": 175 | content["text"] = html.unescape( 176 | chevron.render(content["text"], variables_with_defaults) 177 | ) 178 | 179 | formatted_messages.append(formatted_message) 180 | 181 | return formatted_messages 182 | 183 | @deprecated('Please use "format_messages" instead') 184 | def format(self, variables: Optional[Dict[str, Any]] = None) -> List[Any]: 185 | """ 186 | Deprecated. Please use `format_messages` instead. 187 | """ 188 | return self.format_messages(**(variables or {})) 189 | 190 | def to_langchain_chat_prompt_template(self, additional_messages=[]): 191 | """ 192 | Converts a Literal AI prompt to a LangChain prompt template format. 193 | """ 194 | try: 195 | version("langchain") 196 | except Exception: 197 | raise Exception( 198 | "Please install langchain to use the langchain callback. " 199 | "You can install it with `pip install langchain`" 200 | ) 201 | 202 | from langchain_core.messages import ( 203 | AIMessage, 204 | BaseMessage, 205 | HumanMessage, 206 | SystemMessage, 207 | ) 208 | from langchain_core.prompts import ( 209 | AIMessagePromptTemplate, 210 | ChatPromptTemplate, 211 | HumanMessagePromptTemplate, 212 | SystemMessagePromptTemplate, 213 | ) 214 | 215 | class CustomChatPromptTemplate(ChatPromptTemplate): 216 | orig_messages: Optional[List[GenerationMessage]] = Field( 217 | default_factory=lambda: [] 218 | ) 219 | default_vars: Optional[Dict] = Field(default_factory=lambda: {}) 220 | prompt_id: Optional[str] = None 221 | 222 | def format_messages(self, **kwargs: Any) -> List[BaseMessage]: 223 | variables_with_defaults = { 224 | **(self.default_vars or {}), 225 | **(kwargs or {}), 226 | } 227 | 228 | rendered_messages: List[BaseMessage] = [] 229 | 230 | for index, message in enumerate(self.messages): 231 | content: str = "" 232 | try: 233 | prompt = getattr(message, "prompt") # type: ignore 234 | content = html.unescape( 235 | chevron.render(prompt.template, variables_with_defaults) 236 | ) 237 | except AttributeError: 238 | for m in ChatPromptTemplate.from_messages( 239 | [message] 240 | ).format_messages(): 241 | rendered_messages.append(m) 242 | continue 243 | 244 | additonal_kwargs = {} 245 | if self.orig_messages and index < len(self.orig_messages): 246 | additonal_kwargs = { 247 | "uuid": ( 248 | self.orig_messages[index].get("uuid") 249 | if self.orig_messages 250 | else None 251 | ), 252 | "prompt_id": self.prompt_id, 253 | "variables": variables_with_defaults, 254 | } 255 | 256 | if isinstance(message, HumanMessagePromptTemplate): 257 | rendered_messages.append( 258 | HumanMessage( 259 | content=content, additional_kwargs=additonal_kwargs 260 | ) 261 | ) 262 | if isinstance(message, AIMessagePromptTemplate): 263 | rendered_messages.append( 264 | AIMessage( 265 | content=content, additional_kwargs=additonal_kwargs 266 | ) 267 | ) 268 | if isinstance(message, SystemMessagePromptTemplate): 269 | rendered_messages.append( 270 | SystemMessage( 271 | content=content, additional_kwargs=additonal_kwargs 272 | ) 273 | ) 274 | 275 | return rendered_messages 276 | 277 | async def aformat_messages(self, **kwargs: Any) -> List[BaseMessage]: 278 | return self.format_messages(**kwargs) 279 | 280 | lc_messages = [(m["role"], m["content"]) for m in self.template_messages] 281 | 282 | chat_template = CustomChatPromptTemplate.from_messages( 283 | lc_messages + additional_messages 284 | ) 285 | chat_template.default_vars = self.variables_default_values 286 | chat_template.orig_messages = self.template_messages 287 | chat_template.prompt_id = self.id 288 | 289 | return chat_template 290 | -------------------------------------------------------------------------------- /literalai/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/literalai/py.typed -------------------------------------------------------------------------------- /literalai/requirements.py: -------------------------------------------------------------------------------- 1 | from importlib.metadata import version 2 | 3 | from packaging.requirements import Requirement 4 | 5 | 6 | # Function to check if all packages meet the specified requirements 7 | def check_all_requirements(requirements): 8 | for req_str in requirements: 9 | # Parse the requirement string using packaging.requirements.Requirement 10 | req = Requirement(req_str) 11 | 12 | try: 13 | # Get the installed version of the package 14 | installed_version = version(req.name) 15 | except Exception: 16 | # Package not installed, return False 17 | return False 18 | 19 | # Check if the installed version satisfies the requirement 20 | if not req.specifier.contains(installed_version, prereleases=False): 21 | # Requirement not met, return False 22 | return False 23 | 24 | # All requirements were met, return True 25 | return True 26 | -------------------------------------------------------------------------------- /literalai/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.1.300" 2 | -------------------------------------------------------------------------------- /literalai/wrappers.py: -------------------------------------------------------------------------------- 1 | import time 2 | from functools import wraps 3 | from importlib import import_module 4 | from typing import TYPE_CHECKING, Callable, Optional, TypedDict, Union 5 | 6 | from literalai.context import active_steps_var 7 | 8 | if TYPE_CHECKING: 9 | from literalai.observability.generation import ChatGeneration, CompletionGeneration 10 | from literalai.observability.step import Step 11 | 12 | 13 | class BeforeContext(TypedDict): 14 | original_func: Callable 15 | generation: Optional[Union["ChatGeneration", "CompletionGeneration"]] 16 | step: Optional["Step"] 17 | start: float 18 | 19 | 20 | class AfterContext(TypedDict): 21 | original_func: Callable 22 | generation: Optional[Union["ChatGeneration", "CompletionGeneration"]] 23 | step: Optional["Step"] 24 | start: float 25 | 26 | 27 | def remove_literalai_args(kargs): 28 | """Remove argument prefixed with "literalai_" from kwargs and return them in a separate dict""" 29 | largs = {} 30 | for key in list(kargs.keys()): 31 | if key.startswith("literalai_"): 32 | value = kargs.pop(key) 33 | largs[key] = value 34 | return largs 35 | 36 | 37 | def restore_literalai_args(kargs, largs): 38 | """Reverse the effect of remove_literalai_args by merging the literal arguments into kwargs""" 39 | for key in list(largs.keys()): 40 | kargs[key] = largs[key] 41 | 42 | 43 | def sync_wrapper(before_func=None, after_func=None): 44 | def decorator(original_func): 45 | @wraps(original_func) 46 | def wrapped(*args, **kwargs): 47 | context = {"original_func": original_func} 48 | # If a before_func is provided, call it with the shared context. 49 | if before_func: 50 | before_func(context, *args, **kwargs) 51 | # Remove literal arguments before calling the original function 52 | literalai_args = remove_literalai_args(kwargs) 53 | context["start"] = time.time() 54 | try: 55 | result = original_func(*args, **kwargs) 56 | except Exception as e: 57 | active_steps = active_steps_var.get() 58 | if active_steps and len(active_steps) > 0: 59 | current_step = active_steps[-1] 60 | current_step.error = str(e) 61 | current_step.end() 62 | current_step.processor.flush() 63 | raise e 64 | # If an after_func is provided, call it with the result and the shared context. 65 | if after_func: 66 | restore_literalai_args(kwargs, literalai_args) 67 | result = after_func(result, context, *args, **kwargs) 68 | 69 | return result 70 | 71 | return wrapped 72 | 73 | return decorator 74 | 75 | 76 | def async_wrapper(before_func=None, after_func=None): 77 | def decorator(original_func): 78 | @wraps(original_func) 79 | async def wrapped(*args, **kwargs): 80 | context = {"original_func": original_func} 81 | # If a before_func is provided, call it with the shared context. 82 | if before_func: 83 | await before_func(context, *args, **kwargs) 84 | 85 | # Remove literal arguments before calling the original function 86 | literalai_args = remove_literalai_args(kwargs) 87 | context["start"] = time.time() 88 | try: 89 | result = await original_func(*args, **kwargs) 90 | except Exception as e: 91 | active_steps = active_steps_var.get() 92 | if active_steps and len(active_steps) > 0: 93 | current_step = active_steps[-1] 94 | current_step.error = str(e) 95 | current_step.end() 96 | current_step.processor.flush() 97 | raise e 98 | 99 | # If an after_func is provided, call it with the result and the shared context. 100 | if after_func: 101 | restore_literalai_args(kwargs, literalai_args) 102 | result = await after_func(result, context, *args, **kwargs) 103 | 104 | return result 105 | 106 | return wrapped 107 | 108 | return decorator 109 | 110 | 111 | def wrap_all( 112 | to_wrap: list, 113 | before_wrapper, 114 | after_wrapper, 115 | async_before_wrapper, 116 | async_after_wrapper, 117 | ): 118 | for patch in to_wrap: 119 | module = import_module(str(patch["module"])) 120 | target_object = getattr(module, str(patch["object"])) 121 | original_method = getattr(target_object, str(patch["method"])) 122 | 123 | if patch["async"]: 124 | wrapped_method = async_wrapper( 125 | before_func=async_before_wrapper(metadata=patch["metadata"]), 126 | after_func=async_after_wrapper(metadata=patch["metadata"]), 127 | )(original_method) 128 | else: 129 | wrapped_method = sync_wrapper( 130 | before_func=before_wrapper(metadata=patch["metadata"]), 131 | after_func=after_wrapper(metadata=patch["metadata"]), 132 | )(original_method) 133 | 134 | setattr(target_object, str(patch["method"]), wrapped_method) 135 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | plugins = pydantic.mypy 3 | 4 | [mypy-setuptools.*] 5 | ignore_missing_imports = True 6 | 7 | [mypy-chevron.*] 8 | ignore_missing_imports = True 9 | 10 | [mypy-agents.*] 11 | ignore_missing_imports = True 12 | 13 | [mypy-langchain_community.*] 14 | ignore_missing_imports = True 15 | 16 | [mypy-traceloop.*] 17 | ignore_missing_imports = True 18 | 19 | [mypy-llama_index.*] 20 | ignore_missing_imports = True 21 | 22 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | asyncio_mode=auto 3 | markers = 4 | e2e: end to end test 5 | addopts = 6 | -m 'not e2e' 7 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | pytest-asyncio 3 | pytest-timeout 4 | pytest_httpx==0.30.0 5 | pre-commit 6 | python-dotenv 7 | ruff 8 | mypy 9 | langchain 10 | llama-index 11 | mistralai 12 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | asyncio==3.4.3 2 | packaging==23.2 3 | httpx>=0.23.0 4 | pydantic>=1,<3 5 | openai>=1.0.0 6 | chevron>=0.14.0 7 | traceloop-sdk>=0.33.12 8 | -------------------------------------------------------------------------------- /run-test.sh: -------------------------------------------------------------------------------- 1 | LITERAL_API_URL=http://localhost:3000 LITERAL_API_KEY=my-initial-api-key pytest -m e2e -s -v tests/e2e/ tests/unit/ 2 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import find_packages, setup 2 | 3 | setup( 4 | name="literalai", 5 | version="0.1.300", # update version in literalai/version.py 6 | description="An SDK for observability in Python applications", 7 | long_description=open("README.md").read(), 8 | long_description_content_type="text/markdown", 9 | author="Literal AI", 10 | author_email="contact@literalai.com", 11 | package_data={"literalai": ["py.typed"]}, 12 | packages=find_packages(), 13 | license="Apache License 2.0", 14 | install_requires=[ 15 | "packaging>=23.0", 16 | "httpx>=0.23.0", 17 | "pydantic>=1,<3", 18 | "chevron>=0.14.0", 19 | "traceloop-sdk>=0.33.12", 20 | ], 21 | ) 22 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/tests/__init__.py -------------------------------------------------------------------------------- /tests/e2e/test_llamaindex.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib.parse 3 | 4 | import pytest 5 | from dotenv import load_dotenv 6 | 7 | from literalai import LiteralClient 8 | 9 | load_dotenv() 10 | 11 | 12 | @pytest.fixture 13 | def non_mocked_hosts() -> list: 14 | non_mocked_hosts = [] 15 | 16 | # Always skip mocking API 17 | url = os.getenv("LITERAL_API_URL", None) 18 | if url is not None: 19 | parsed = urllib.parse.urlparse(url) 20 | non_mocked_hosts.append(parsed.hostname) 21 | 22 | return non_mocked_hosts 23 | 24 | 25 | @pytest.mark.e2e 26 | class TestLlamaIndex: 27 | @pytest.fixture( 28 | scope="class" 29 | ) # Feel free to move this fixture up for further testing 30 | def client(self): 31 | url = os.getenv("LITERAL_API_URL", None) 32 | api_key = os.getenv("LITERAL_API_KEY", None) 33 | assert url is not None and api_key is not None, "Missing environment variables" 34 | 35 | client = LiteralClient(batch_size=5, url=url, api_key=api_key) 36 | client.instrument_llamaindex() 37 | 38 | return client 39 | 40 | async def test_instrument_llamaindex(self, client: "LiteralClient"): 41 | assert client is not None 42 | -------------------------------------------------------------------------------- /tests/e2e/test_mistralai.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib.parse 3 | from asyncio import sleep 4 | 5 | import pytest 6 | from mistralai import Mistral 7 | from pytest_httpx import HTTPXMock 8 | 9 | from literalai.client import LiteralClient 10 | from literalai.observability.generation import ChatGeneration, CompletionGeneration 11 | 12 | 13 | @pytest.fixture 14 | def non_mocked_hosts() -> list: 15 | non_mocked_hosts = [] 16 | 17 | # Always skip mocking API 18 | url = os.getenv("LITERAL_API_URL", None) 19 | if url is not None: 20 | parsed = urllib.parse.urlparse(url) 21 | non_mocked_hosts.append(parsed.hostname) 22 | 23 | return non_mocked_hosts 24 | 25 | 26 | @pytest.mark.e2e 27 | class TestMistralAI: 28 | @pytest.fixture(scope="class") 29 | def client(self): 30 | url = os.getenv("LITERAL_API_URL", None) 31 | api_key = os.getenv("LITERAL_API_KEY", None) 32 | assert url is not None and api_key is not None, "Missing environment variables" 33 | 34 | client = LiteralClient(batch_size=5, url=url, api_key=api_key) 35 | client.instrument_mistralai() 36 | 37 | return client 38 | 39 | async def test_chat(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 40 | httpx_mock.add_response( 41 | json={ 42 | "id": "afc02d747c9b47d3b21e6a1f9fd7e2e6", 43 | "object": "chat.completion", 44 | "created": 1718881182, 45 | "model": "open-mistral-7b", 46 | "choices": [ 47 | { 48 | "index": 0, 49 | "message": { 50 | "role": "assistant", 51 | "content": '1+1=2\n\nHere\'s a fun fact: The sum of 1 and 1 is called the "successor" in mathematical logic and set theory. It represents the next number after 1. This concept is fundamental in the development of mathematics and forms the basis for the understanding of numbers, arithmetic, and more complex mathematical structures.', 52 | "tool_calls": None, 53 | }, 54 | "finish_reason": "stop", 55 | "logprobs": None, 56 | } 57 | ], 58 | "usage": { 59 | "prompt_tokens": 9, 60 | "total_tokens": 85, 61 | "completion_tokens": 76, 62 | }, 63 | } 64 | ) 65 | mai_client = Mistral(api_key="j3s4V1z4") 66 | thread_id = None 67 | 68 | @client.thread 69 | def main(): 70 | # https://docs.mistral.ai/api/#operation/createChatCompletion 71 | mai_client.chat.complete( 72 | model="open-mistral-7b", 73 | messages=[ 74 | { 75 | "role": "user", 76 | "content": "1+1=?", 77 | } 78 | ], 79 | temperature=0, 80 | max_tokens=256, 81 | ) 82 | return client.get_current_thread() 83 | 84 | thread_id = main().id 85 | await sleep(2) 86 | thread = client.api.get_thread(id=thread_id) 87 | assert thread is not None 88 | assert thread.steps is not None 89 | assert len(thread.steps) == 1 90 | 91 | step = thread.steps[0] 92 | 93 | assert step.type == "llm" 94 | assert step.generation is not None 95 | assert type(step.generation) is ChatGeneration 96 | assert step.generation.settings is not None 97 | assert step.generation.model == "open-mistral-7b" 98 | 99 | async def test_completion(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 100 | httpx_mock.add_response( 101 | json={ 102 | "id": "8103af31e335493da136a79f2e64b59c", 103 | "object": "chat.completion", 104 | "created": 1718884349, 105 | "model": "codestral-2405", 106 | "choices": [ 107 | { 108 | "index": 0, 109 | "message": { 110 | "role": "assistant", 111 | "content": "2\n\n", 112 | "tool_calls": None, 113 | }, 114 | "finish_reason": "length", 115 | "logprobs": None, 116 | } 117 | ], 118 | "usage": { 119 | "prompt_tokens": 8, 120 | "total_tokens": 11, 121 | "completion_tokens": 3, 122 | }, 123 | }, 124 | ) 125 | 126 | mai_client = Mistral(api_key="j3s4V1z4") 127 | thread_id = None 128 | 129 | @client.thread 130 | def main(): 131 | # https://docs.mistral.ai/api/#operation/createFIMCompletion 132 | mai_client.fim.complete( 133 | model="codestral-2405", 134 | prompt="1+1=", 135 | temperature=0, 136 | max_tokens=3, 137 | ) 138 | return client.get_current_thread() 139 | 140 | thread_id = main().id 141 | await sleep(2) 142 | thread = client.api.get_thread(id=thread_id) 143 | assert thread is not None 144 | assert thread.steps is not None 145 | assert len(thread.steps) == 1 146 | 147 | step = thread.steps[0] 148 | 149 | assert step.type == "llm" 150 | assert step.generation is not None 151 | assert type(step.generation) is CompletionGeneration 152 | assert step.generation.settings is not None 153 | assert step.generation.model == "codestral-2405" 154 | assert step.generation.completion == "2\n\n" 155 | assert step.generation.token_count == 11 156 | assert step.generation.prompt == "1+1=" 157 | 158 | async def test_async_chat(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 159 | httpx_mock.add_response( 160 | json={ 161 | "id": "afc02d747c9b47d3b21e6a1f9fd7e2e6", 162 | "object": "chat.completion", 163 | "created": 1718881182, 164 | "model": "open-mistral-7b", 165 | "choices": [ 166 | { 167 | "index": 0, 168 | "message": { 169 | "role": "assistant", 170 | "content": '1+1=2\n\nHere\'s a fun fact: The sum of 1 and 1 is called the "successor" in mathematical logic and set theory. It represents the next number after 1. This concept is fundamental in the development of mathematics and forms the basis for the understanding of numbers, arithmetic, and more complex mathematical structures.', 171 | "tool_calls": None, 172 | }, 173 | "finish_reason": "stop", 174 | "logprobs": None, 175 | } 176 | ], 177 | "usage": { 178 | "prompt_tokens": 9, 179 | "total_tokens": 85, 180 | "completion_tokens": 76, 181 | }, 182 | }, 183 | ) 184 | 185 | mai_client = Mistral(api_key="j3s4V1z4") 186 | thread_id = None 187 | 188 | @client.thread 189 | async def main(): 190 | # https://docs.mistral.ai/api/#operation/createChatCompletion 191 | await mai_client.chat.complete_async( 192 | model="open-mistral-7b", 193 | messages=[ 194 | { 195 | "role": "user", 196 | "content": "1+1=?", 197 | } 198 | ], 199 | temperature=0, 200 | max_tokens=256, 201 | ) 202 | return client.get_current_thread() 203 | 204 | thread_id = (await main()).id 205 | await sleep(2) 206 | thread = client.api.get_thread(id=thread_id) 207 | assert thread is not None 208 | assert thread.steps is not None 209 | assert len(thread.steps) == 1 210 | 211 | step = thread.steps[0] 212 | 213 | assert step.type == "llm" 214 | assert step.generation is not None 215 | assert type(step.generation) is ChatGeneration 216 | assert step.generation.settings is not None 217 | assert step.generation.model == "open-mistral-7b" 218 | 219 | async def test_async_completion( 220 | self, client: "LiteralClient", httpx_mock: "HTTPXMock" 221 | ): 222 | httpx_mock.add_response( 223 | json={ 224 | "id": "8103af31e335493da136a79f2e64b59c", 225 | "object": "chat.completion", 226 | "created": 1718884349, 227 | "model": "codestral-2405", 228 | "choices": [ 229 | { 230 | "index": 0, 231 | "message": { 232 | "role": "assistant", 233 | "content": "2\n\n", 234 | "tool_calls": None, 235 | }, 236 | "finish_reason": "length", 237 | "logprobs": None, 238 | } 239 | ], 240 | "usage": { 241 | "prompt_tokens": 8, 242 | "total_tokens": 11, 243 | "completion_tokens": 3, 244 | }, 245 | }, 246 | ) 247 | 248 | mai_client = Mistral(api_key="j3s4V1z4") 249 | thread_id = None 250 | 251 | @client.thread 252 | async def main(): 253 | # https://docs.mistral.ai/api/#operation/createFIMCompletion 254 | await mai_client.fim.complete_async( 255 | model="codestral-2405", 256 | prompt="1+1=", 257 | temperature=0, 258 | max_tokens=3, 259 | ) 260 | return client.get_current_thread() 261 | 262 | thread_id = (await main()).id 263 | await sleep(2) 264 | thread = client.api.get_thread(id=thread_id) 265 | assert thread is not None 266 | assert thread.steps is not None 267 | assert len(thread.steps) == 1 268 | 269 | step = thread.steps[0] 270 | 271 | assert step.type == "llm" 272 | assert step.generation is not None 273 | assert type(step.generation) is CompletionGeneration 274 | assert step.generation.settings is not None 275 | assert step.generation.model == "codestral-2405" 276 | assert step.generation.completion == "2\n\n" 277 | assert step.generation.token_count == 11 278 | assert step.generation.prompt == "1+1=" 279 | -------------------------------------------------------------------------------- /tests/e2e/test_openai.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib.parse 3 | from asyncio import sleep 4 | 5 | import pytest 6 | from openai import AsyncOpenAI, AzureOpenAI, OpenAI 7 | from pytest_httpx import HTTPXMock 8 | 9 | from literalai import LiteralClient 10 | from literalai.observability.generation import ChatGeneration, CompletionGeneration 11 | 12 | 13 | @pytest.fixture 14 | def non_mocked_hosts() -> list: 15 | non_mocked_hosts = [] 16 | 17 | # Always skip mocking API 18 | url = os.getenv("LITERAL_API_URL", None) 19 | if url is not None: 20 | parsed = urllib.parse.urlparse(url) 21 | non_mocked_hosts.append(parsed.hostname) 22 | 23 | return non_mocked_hosts 24 | 25 | 26 | @pytest.mark.e2e 27 | class TestOpenAI: 28 | @pytest.fixture( 29 | scope="class" 30 | ) # Feel free to move this fixture up for further testing 31 | def client(self): 32 | url = os.getenv("LITERAL_API_URL", None) 33 | api_key = os.getenv("LITERAL_API_KEY", None) 34 | assert url is not None and api_key is not None, "Missing environment variables" 35 | 36 | client = LiteralClient(batch_size=5, url=url, api_key=api_key) 37 | client.instrument_openai() 38 | 39 | return client 40 | 41 | async def test_chat(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 42 | # https://platform.openai.com/docs/api-reference/chat/object 43 | httpx_mock.add_response( 44 | json={ 45 | "id": "chatcmpl-123", 46 | "object": "chat.completion", 47 | "created": 1677652288, 48 | "model": "gpt-3.5-turbo-0613", 49 | "system_fingerprint": "fp_44709d6fcb", 50 | "choices": [ 51 | { 52 | "index": 0, 53 | "message": { 54 | "role": "assistant", 55 | "content": "\n\nHello there, how may I assist you today?", 56 | }, 57 | "logprobs": None, 58 | "finish_reason": "stop", 59 | } 60 | ], 61 | "usage": { 62 | "prompt_tokens": 9, 63 | "completion_tokens": 12, 64 | "total_tokens": 21, 65 | }, 66 | }, 67 | ) 68 | 69 | openai_client = OpenAI(api_key="sk_test_123") 70 | thread_id = None 71 | 72 | @client.thread 73 | def main(): 74 | # https://platform.openai.com/docs/api-reference/chat/create 75 | openai_client.chat.completions.create( 76 | model="gpt-3.5-turbo", 77 | messages=[ 78 | { 79 | "role": "user", 80 | "content": "Tell me a funny joke.", 81 | } 82 | ], 83 | temperature=1, 84 | max_tokens=256, 85 | top_p=1, 86 | frequency_penalty=0, 87 | presence_penalty=0, 88 | ) 89 | return client.get_current_thread() 90 | 91 | thread_id = main().id 92 | await sleep(2) 93 | thread = client.api.get_thread(id=thread_id) 94 | assert thread is not None 95 | assert thread.steps is not None 96 | assert len(thread.steps) == 1 97 | 98 | step = thread.steps[0] 99 | 100 | assert step.type == "llm" 101 | assert step.generation is not None 102 | assert type(step.generation) is ChatGeneration 103 | assert step.generation.settings is not None 104 | assert step.generation.model == "gpt-3.5-turbo-0613" 105 | 106 | async def test_completion(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 107 | # https://platform.openai.com/docs/api-reference/completions/object 108 | httpx_mock.add_response( 109 | json={ 110 | "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7", 111 | "object": "text_completion", 112 | "created": 1589478378, 113 | "model": "gpt-3.5-turbo", 114 | "choices": [ 115 | { 116 | "text": "\n\nThis is indeed a test", 117 | "index": 0, 118 | "logprobs": None, 119 | "finish_reason": "length", 120 | } 121 | ], 122 | "usage": { 123 | "prompt_tokens": 5, 124 | "completion_tokens": 7, 125 | "total_tokens": 12, 126 | }, 127 | }, 128 | ) 129 | 130 | openai_client = OpenAI(api_key="sk_test_123") 131 | thread_id = None 132 | 133 | @client.thread 134 | def main(): 135 | # https://platform.openai.com/docs/api-reference/completions/create 136 | openai_client.completions.create( 137 | model="gpt-3.5-turbo", 138 | prompt="Tell me a funny joke.", 139 | temperature=1, 140 | max_tokens=256, 141 | top_p=1, 142 | frequency_penalty=0, 143 | presence_penalty=0, 144 | ) 145 | return client.get_current_thread() 146 | 147 | thread_id = main().id 148 | await sleep(2) 149 | thread = client.api.get_thread(id=thread_id) 150 | assert thread is not None 151 | assert thread.steps is not None 152 | assert len(thread.steps) == 1 153 | 154 | step = thread.steps[0] 155 | 156 | assert step.type == "llm" 157 | assert step.generation is not None 158 | assert type(step.generation) is CompletionGeneration 159 | assert step.generation.settings is not None 160 | assert step.generation.model == "gpt-3.5-turbo" 161 | assert step.generation.completion == "\n\nThis is indeed a test" 162 | assert step.generation.token_count == 12 163 | assert step.generation.prompt == "Tell me a funny joke." 164 | 165 | async def test_async_chat(self, client: "LiteralClient", httpx_mock: "HTTPXMock"): 166 | # https://platform.openai.com/docs/api-reference/chat/object 167 | httpx_mock.add_response( 168 | json={ 169 | "id": "chatcmpl-123", 170 | "object": "chat.completion", 171 | "created": 1677652288, 172 | "model": "gpt-3.5-turbo-0613", 173 | "system_fingerprint": "fp_44709d6fcb", 174 | "choices": [ 175 | { 176 | "index": 0, 177 | "message": { 178 | "role": "assistant", 179 | "content": "\n\nHello there, how may I assist you today?", 180 | }, 181 | "logprobs": None, 182 | "finish_reason": "stop", 183 | } 184 | ], 185 | "usage": { 186 | "prompt_tokens": 9, 187 | "completion_tokens": 12, 188 | "total_tokens": 21, 189 | }, 190 | }, 191 | ) 192 | 193 | openai_client = AsyncOpenAI(api_key="sk_test_123") 194 | thread_id = None 195 | 196 | @client.thread 197 | async def main(): 198 | # https://platform.openai.com/docs/api-reference/chat/create 199 | await openai_client.chat.completions.create( 200 | model="gpt-3.5-turbo", 201 | messages=[ 202 | { 203 | "role": "user", 204 | "content": "Tell me a funny joke.", 205 | } 206 | ], 207 | temperature=1, 208 | max_tokens=256, 209 | top_p=1, 210 | frequency_penalty=0, 211 | presence_penalty=0, 212 | ) 213 | return client.get_current_thread() 214 | 215 | thread_id = (await main()).id 216 | await sleep(2) 217 | thread = client.api.get_thread(id=thread_id) 218 | assert thread is not None 219 | assert thread.steps is not None 220 | assert len(thread.steps) == 1 221 | 222 | step = thread.steps[0] 223 | 224 | assert step.type == "llm" 225 | assert step.generation is not None 226 | assert type(step.generation) is ChatGeneration 227 | assert step.generation.settings is not None 228 | assert step.generation.model == "gpt-3.5-turbo-0613" 229 | 230 | async def test_async_completion( 231 | self, client: "LiteralClient", httpx_mock: "HTTPXMock" 232 | ): 233 | # https://platform.openai.com/docs/api-reference/completions/object 234 | httpx_mock.add_response( 235 | json={ 236 | "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7", 237 | "object": "text_completion", 238 | "created": 1589478378, 239 | "model": "gpt-3.5-turbo", 240 | "choices": [ 241 | { 242 | "text": "\n\nThis is indeed a test", 243 | "index": 0, 244 | "logprobs": None, 245 | "finish_reason": "length", 246 | } 247 | ], 248 | "usage": { 249 | "prompt_tokens": 5, 250 | "completion_tokens": 7, 251 | "total_tokens": 12, 252 | }, 253 | }, 254 | ) 255 | 256 | openai_client = AsyncOpenAI(api_key="sk_test_123") 257 | thread_id = None 258 | 259 | @client.thread 260 | async def main(): 261 | # https://platform.openai.com/docs/api-reference/completions/create 262 | await openai_client.completions.create( 263 | model="gpt-3.5-turbo", 264 | prompt="Tell me a funny joke.", 265 | temperature=1, 266 | max_tokens=256, 267 | top_p=1, 268 | frequency_penalty=0, 269 | presence_penalty=0, 270 | ) 271 | return client.get_current_thread() 272 | 273 | thread_id = (await main()).id 274 | await sleep(2) 275 | thread = client.api.get_thread(id=thread_id) 276 | assert thread is not None 277 | assert thread.steps is not None 278 | assert len(thread.steps) == 1 279 | 280 | step = thread.steps[0] 281 | 282 | assert step.type == "llm" 283 | assert step.generation is not None 284 | assert type(step.generation) is CompletionGeneration 285 | assert step.generation.settings is not None 286 | assert step.generation.model == "gpt-3.5-turbo" 287 | assert step.generation.completion == "\n\nThis is indeed a test" 288 | assert step.generation.token_count == 12 289 | assert step.generation.prompt == "Tell me a funny joke." 290 | 291 | async def test_azure_completion( 292 | self, client: "LiteralClient", httpx_mock: "HTTPXMock" 293 | ): 294 | # https://platform.openai.com/docs/api-reference/completions/object 295 | httpx_mock.add_response( 296 | json={ 297 | "id": "cmpl-uqkvlQyYK7bGYrRHQ0eXlWi7", 298 | "object": "text_completion", 299 | "created": 1589478378, 300 | "model": "gpt-3.5-turbo", 301 | "choices": [ 302 | { 303 | "text": "\n\nThis is indeed a test", 304 | "index": 0, 305 | "logprobs": None, 306 | "finish_reason": "length", 307 | } 308 | ], 309 | "usage": { 310 | "prompt_tokens": 5, 311 | "completion_tokens": 7, 312 | "total_tokens": 12, 313 | }, 314 | }, 315 | ) 316 | 317 | openai_client = AzureOpenAI( 318 | api_key="sk_test_123", 319 | api_version="2023-05-15", 320 | azure_endpoint="https://example.org", 321 | ) 322 | thread_id = None 323 | 324 | @client.thread 325 | def main(): 326 | # https://platform.openai.com/docs/api-reference/completions/create 327 | openai_client.completions.create( 328 | model="gpt-3.5-turbo", 329 | prompt="Tell me a funny joke.", 330 | temperature=1, 331 | max_tokens=256, 332 | top_p=1, 333 | frequency_penalty=0, 334 | presence_penalty=0, 335 | ) 336 | return client.get_current_thread() 337 | 338 | thread_id = main().id 339 | await sleep(2) 340 | thread = client.api.get_thread(id=thread_id) 341 | assert thread is not None 342 | assert thread.steps is not None 343 | assert len(thread.steps) == 1 344 | 345 | step = thread.steps[0] 346 | 347 | assert step.type == "llm" 348 | assert step.generation is not None 349 | assert type(step.generation) is CompletionGeneration 350 | assert step.generation.settings is not None 351 | assert step.generation.model == "gpt-3.5-turbo" 352 | assert step.generation.completion == "\n\nThis is indeed a test" 353 | assert step.generation.token_count == 12 354 | assert step.generation.prompt == "Tell me a funny joke." 355 | -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Chainlit/literalai-python/914ef9b9402050103aeeb70670ca9f0033bff55f/tests/unit/__init__.py -------------------------------------------------------------------------------- /tests/unit/test_cache.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from literalai.api import LiteralAPI 4 | from literalai.cache.prompt_helpers import put_prompt 5 | from literalai.cache.shared_cache import SharedCache 6 | from literalai.prompt_engineering.prompt import Prompt 7 | 8 | 9 | def default_prompt(id: str = "1", name: str = "test", version: int = 1) -> Prompt: 10 | return Prompt( 11 | api=LiteralAPI(), 12 | id=id, 13 | name=name, 14 | version=version, 15 | created_at="", 16 | updated_at="", 17 | type="chat", # type: ignore 18 | url="", 19 | version_desc=None, 20 | template_messages=[], 21 | tools=None, 22 | provider="", 23 | settings={}, 24 | variables=[], 25 | variables_default_values=None, 26 | ) 27 | 28 | 29 | def test_singleton_instance(): 30 | """Test that SharedCache maintains singleton pattern""" 31 | cache1 = SharedCache() 32 | cache2 = SharedCache() 33 | assert cache1 is cache2 34 | 35 | 36 | def test_get_empty_cache(): 37 | """Test getting from empty cache returns None""" 38 | cache = SharedCache() 39 | cache.clear() 40 | 41 | assert cache.get_cache() == {} 42 | 43 | 44 | def test_put_and_get_prompt_by_id_by_name_version_by_name(): 45 | """Test storing and retrieving prompt by ID by name-version by name""" 46 | cache = SharedCache() 47 | cache.clear() 48 | 49 | prompt = default_prompt() 50 | put_prompt(cache, prompt) 51 | 52 | retrieved_by_id = cache.get(id="1") 53 | assert retrieved_by_id is prompt 54 | 55 | retrieved_by_name_version = cache.get(name="test", version=1) 56 | assert retrieved_by_name_version is prompt 57 | 58 | retrieved_by_name = cache.get(name="test") 59 | assert retrieved_by_name is prompt 60 | 61 | 62 | def test_clear_cache(): 63 | """Test clearing the cache""" 64 | cache = SharedCache() 65 | prompt = default_prompt() 66 | put_prompt(cache, prompt) 67 | 68 | cache.clear() 69 | assert cache.get_cache() == {} 70 | 71 | 72 | def test_update_existing_prompt(): 73 | """Test updating an existing prompt""" 74 | cache = SharedCache() 75 | cache.clear() 76 | 77 | prompt1 = default_prompt() 78 | prompt2 = default_prompt(id="1", version=2) 79 | 80 | cache.put_prompt(prompt1) 81 | cache.put_prompt(prompt2) 82 | 83 | retrieved = cache.get(id="1") 84 | assert retrieved is prompt2 85 | assert retrieved.version == 2 86 | 87 | 88 | def test_error_handling(): 89 | """Test error handling for invalid inputs""" 90 | cache = SharedCache() 91 | cache.clear() 92 | 93 | assert cache.get_cache() == {} 94 | assert cache.get(key="") is None 95 | 96 | with pytest.raises(TypeError): 97 | cache.get(5) # type: ignore 98 | 99 | with pytest.raises(TypeError): 100 | cache.put(5, "test") # type: ignore 101 | --------------------------------------------------------------------------------