├── .env.example ├── .github ├── dependabot.yml └── workflows │ ├── cla.yml │ ├── lint.yml │ ├── release-graphiti-core.yml │ ├── typecheck.yml │ └── unit_tests.yml ├── .gitignore ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── Dockerfile ├── LICENSE ├── Makefile ├── README.md ├── SECURITY.md ├── Zep-CLA.md ├── conftest.py ├── depot.json ├── docker-compose.test.yml ├── docker-compose.yml ├── ellipsis.yaml ├── examples ├── data │ └── manybirds_products.json ├── ecommerce │ ├── runner.ipynb │ └── runner.py ├── langgraph-agent │ ├── agent.ipynb │ └── tinybirds-jess.png ├── podcast │ ├── podcast_runner.py │ ├── podcast_transcript.txt │ └── transcript_parser.py ├── quickstart │ ├── README.md │ ├── quickstart.py │ └── requirements.txt └── wizard_of_oz │ ├── parser.py │ ├── runner.py │ └── woo.txt ├── graphiti_core ├── __init__.py ├── cross_encoder │ ├── __init__.py │ ├── bge_reranker_client.py │ ├── client.py │ └── openai_reranker_client.py ├── edges.py ├── embedder │ ├── __init__.py │ ├── client.py │ ├── gemini.py │ ├── openai.py │ └── voyage.py ├── errors.py ├── graphiti.py ├── graphiti_types.py ├── helpers.py ├── llm_client │ ├── __init__.py │ ├── anthropic_client.py │ ├── client.py │ ├── config.py │ ├── errors.py │ ├── gemini_client.py │ ├── groq_client.py │ ├── openai_client.py │ ├── openai_generic_client.py │ └── utils.py ├── models │ ├── __init__.py │ ├── edges │ │ ├── __init__.py │ │ └── edge_db_queries.py │ └── nodes │ │ ├── __init__.py │ │ └── node_db_queries.py ├── nodes.py ├── prompts │ ├── __init__.py │ ├── dedupe_edges.py │ ├── dedupe_nodes.py │ ├── eval.py │ ├── extract_edge_dates.py │ ├── extract_edges.py │ ├── extract_nodes.py │ ├── invalidate_edges.py │ ├── lib.py │ ├── models.py │ ├── prompt_helpers.py │ └── summarize_nodes.py ├── py.typed ├── search │ ├── __init__.py │ ├── search.py │ ├── search_config.py │ ├── search_config_recipes.py │ ├── search_filters.py │ ├── search_helpers.py │ └── search_utils.py └── utils │ ├── __init__.py │ ├── bulk_utils.py │ ├── datetime_utils.py │ ├── maintenance │ ├── __init__.py │ ├── community_operations.py │ ├── edge_operations.py │ ├── graph_data_operations.py │ ├── node_operations.py │ ├── temporal_operations.py │ └── utils.py │ └── ontology_utils │ └── entity_types_utils.py ├── images ├── arxiv-screenshot.png ├── graphiti-graph-intro.gif ├── graphiti-intro-slides-stock-2.gif └── simple_graph.svg ├── mcp_server ├── .env.example ├── .python-version ├── Dockerfile ├── README.md ├── cursor_rules.md ├── docker-compose.yml ├── graphiti_mcp_server.py ├── mcp_config_sse_example.json ├── mcp_config_stdio_example.json ├── pyproject.toml └── uv.lock ├── poetry.lock ├── py.typed ├── pyproject.toml ├── pytest.ini ├── server ├── .env.example ├── Makefile ├── README.md ├── graph_service │ ├── __init__.py │ ├── config.py │ ├── dto │ │ ├── __init__.py │ │ ├── common.py │ │ ├── ingest.py │ │ └── retrieve.py │ ├── main.py │ ├── routers │ │ ├── __init__.py │ │ ├── ingest.py │ │ └── retrieve.py │ └── zep_graphiti.py ├── poetry.lock └── pyproject.toml ├── signatures └── version1 │ └── cla.json └── tests ├── cross_encoder └── test_bge_reranker_client.py ├── embedder ├── embedder_fixtures.py ├── test_gemini.py ├── test_openai.py └── test_voyage.py ├── evals ├── data │ └── longmemeval_data │ │ ├── README.md │ │ └── longmemeval_oracle.json ├── eval_cli.py ├── eval_e2e_graph_building.py ├── pytest.ini └── utils.py ├── helpers_test.py ├── llm_client ├── test_anthropic_client.py ├── test_anthropic_client_int.py ├── test_client.py └── test_errors.py ├── test_graphiti_int.py ├── test_node_int.py └── utils ├── maintenance ├── test_edge_operations.py └── test_temporal_operations_int.py └── search └── search_utils_test.py /.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY= 2 | NEO4J_URI= 3 | NEO4J_PORT= 4 | NEO4J_USER= 5 | NEO4J_PASSWORD= 6 | DEFAULT_DATABASE= 7 | USE_PARALLEL_RUNTIME= 8 | SEMAPHORE_LIMIT= 9 | GITHUB_SHA= 10 | MAX_REFLEXION_ITERATIONS= 11 | ANTHROPIC_API_KEY= -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | # To get started with Dependabot version updates, you'll need to specify which 2 | # package ecosystems to update and where the package manifests are located. 3 | # Please see the documentation for all configuration options: 4 | # https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file 5 | 6 | version: 2 7 | updates: 8 | - package-ecosystem: "pip" # See documentation for possible values 9 | directory: "/" # Location of package manifests 10 | schedule: 11 | interval: "weekly" 12 | -------------------------------------------------------------------------------- /.github/workflows/cla.yml: -------------------------------------------------------------------------------- 1 | name: "CLA Assistant" 2 | on: 3 | issue_comment: 4 | types: [created] 5 | pull_request_target: 6 | types: [opened, closed, synchronize] 7 | 8 | # explicitly configure permissions, in case your GITHUB_TOKEN workflow permissions are set to read-only in repository settings 9 | permissions: 10 | actions: write 11 | contents: write # this can be 'read' if the signatures are in remote repository 12 | pull-requests: write 13 | statuses: write 14 | 15 | jobs: 16 | CLAAssistant: 17 | runs-on: ubuntu-latest 18 | steps: 19 | - name: "CLA Assistant" 20 | if: (github.event.comment.body == 'recheck' || github.event.comment.body == 'I have read the CLA Document and I hereby sign the CLA') || github.event_name == 'pull_request_target' 21 | uses: contributor-assistant/github-action@v2.6.1 22 | env: 23 | # the default github token does not have branch protection override permissions 24 | # the repo secrets will need to be updated when the token expires. 25 | GITHUB_TOKEN: ${{ secrets.DANIEL_PAT }} 26 | with: 27 | path-to-signatures: "signatures/version1/cla.json" 28 | path-to-document: "https://github.com/getzep/graphiti/blob/main/Zep-CLA.md" # e.g. a CLA or a DCO document 29 | # branch should not be protected unless a personal PAT is used 30 | branch: "main" 31 | allowlist: paul-paliychuk,prasmussen15,danielchalef,dependabot[bot],ellipsis-dev 32 | 33 | # the followings are the optional inputs - If the optional inputs are not given, then default values will be taken 34 | #remote-organization-name: enter the remote organization name where the signatures should be stored (Default is storing the signatures in the same repository) 35 | #remote-repository-name: enter the remote repository name where the signatures should be stored (Default is storing the signatures in the same repository) 36 | #create-file-commit-message: 'For example: Creating file for storing CLA Signatures' 37 | #signed-commit-message: 'For example: $contributorName has signed the CLA in $owner/$repo#$pullRequestNo' 38 | #custom-notsigned-prcomment: 'pull request comment with Introductory message to ask new contributors to sign' 39 | #custom-pr-sign-comment: 'The signature to be committed in order to sign the CLA' 40 | #custom-allsigned-prcomment: 'pull request comment when all contributors has signed, defaults to **CLA Assistant Lite bot** All Contributors have signed the CLA.' 41 | #lock-pullrequest-aftermerge: false - if you don't want this bot to automatically lock the pull request after merging (default - true) 42 | #use-dco-flag: true - If you are using DCO instead of CLA 43 | -------------------------------------------------------------------------------- /.github/workflows/lint.yml: -------------------------------------------------------------------------------- 1 | name: Lint with Ruff 2 | 3 | on: 4 | push: 5 | branches: ["main"] 6 | pull_request: 7 | branches: ["main"] 8 | 9 | jobs: 10 | ruff: 11 | environment: development 12 | runs-on: depot-ubuntu-22.04 13 | steps: 14 | - uses: actions/checkout@v4 15 | - name: Set up Python 16 | uses: actions/setup-python@v5 17 | with: 18 | python-version: "3.10" 19 | - name: Install dependencies 20 | run: | 21 | python -m pip install --upgrade pip 22 | pip install "ruff>0.1.7" 23 | - name: Run Ruff linting 24 | run: ruff check --output-format=github 25 | -------------------------------------------------------------------------------- /.github/workflows/release-graphiti-core.yml: -------------------------------------------------------------------------------- 1 | name: Release to PyPI 2 | 3 | on: 4 | push: 5 | tags: ["v*.*.*"] 6 | 7 | env: 8 | POETRY_VERSION: "2.1.2" 9 | 10 | jobs: 11 | release: 12 | runs-on: ubuntu-latest 13 | permissions: 14 | id-token: write 15 | contents: write 16 | environment: 17 | name: release 18 | url: https://pypi.org/p/zep-cloud 19 | steps: 20 | - uses: actions/checkout@v4 21 | - name: Install poetry 22 | run: pipx install poetry==$POETRY_VERSION 23 | - name: Set up Python 3.10 24 | uses: actions/setup-python@v5 25 | with: 26 | python-version: "3.10" 27 | cache: "poetry" 28 | - name: Compare pyproject version with tag 29 | run: | 30 | TAG_VERSION=${GITHUB_REF#refs/tags/} 31 | if [ "$TAG_VERSION" != "v$(poetry version --short)" ]; then 32 | echo "Tag version $TAG_VERSION does not match the project version $(poetry version --short)" 33 | exit 1 34 | fi 35 | - name: Build project for distribution 36 | run: poetry build 37 | - name: Publish package distributions to PyPI 38 | uses: pypa/gh-action-pypi-publish@release/v1 39 | -------------------------------------------------------------------------------- /.github/workflows/typecheck.yml: -------------------------------------------------------------------------------- 1 | name: MyPy Type Check 2 | 3 | on: 4 | push: 5 | branches: ["main"] 6 | pull_request: 7 | branches: ["main"] 8 | 9 | jobs: 10 | mypy: 11 | runs-on: depot-ubuntu-22.04 12 | environment: development 13 | steps: 14 | - uses: actions/checkout@v4 15 | - name: Set up Python 16 | id: setup-python 17 | uses: actions/setup-python@v5 18 | with: 19 | python-version: "3.10" 20 | - name: Install Poetry 21 | uses: snok/install-poetry@v1 22 | with: 23 | version: 2.1.2 24 | virtualenvs-create: true 25 | virtualenvs-in-project: true 26 | installer-parallel: true 27 | - name: Install dependencies 28 | run: poetry install --no-interaction --with dev 29 | - name: Run MyPy for graphiti-core 30 | shell: bash 31 | run: | 32 | set -o pipefail 33 | poetry run mypy ./graphiti_core --show-column-numbers --show-error-codes | sed -E ' 34 | s/^(.*):([0-9]+):([0-9]+): (error|warning): (.+) \[(.+)\]/::error file=\1,line=\2,endLine=\2,col=\3,title=\6::\5/; 35 | s/^(.*):([0-9]+):([0-9]+): note: (.+)/::notice file=\1,line=\2,endLine=\2,col=\3,title=Note::\4/; 36 | ' 37 | - name: Install graph-service dependencies 38 | shell: bash 39 | run: | 40 | cd server 41 | poetry install --no-interaction --with dev 42 | - name: Run MyPy for graph-service 43 | shell: bash 44 | run: | 45 | cd server 46 | set -o pipefail 47 | poetry run mypy . --show-column-numbers --show-error-codes | sed -E ' 48 | s/^(.*):([0-9]+):([0-9]+): (error|warning): (.+) \[(.+)\]/::error file=\1,line=\2,endLine=\2,col=\3,title=\6::\5/; 49 | s/^(.*):([0-9]+):([0-9]+): note: (.+)/::notice file=\1,line=\2,endLine=\2,col=\3,title=Note::\4/; 50 | ' 51 | -------------------------------------------------------------------------------- /.github/workflows/unit_tests.yml: -------------------------------------------------------------------------------- 1 | name: Unit Tests 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | pull_request: 7 | branches: [main] 8 | 9 | jobs: 10 | test: 11 | runs-on: depot-ubuntu-22.04 12 | environment: 13 | name: development 14 | steps: 15 | - uses: actions/checkout@v4 16 | - name: Set up Python 17 | uses: actions/setup-python@v5 18 | with: 19 | python-version: "3.10" 20 | - name: Load cached Poetry installation 21 | uses: actions/cache@v4 22 | with: 23 | path: ~/.local 24 | key: poetry-${{ runner.os }}-${{ hashFiles('**/poetry.lock') }} 25 | - name: Install Poetry 26 | uses: snok/install-poetry@v1 27 | with: 28 | version: 2.1.2 29 | virtualenvs-create: true 30 | virtualenvs-in-project: true 31 | installer-parallel: true 32 | - name: Install dependencies 33 | run: poetry install --no-interaction --no-root 34 | - name: Run non-integration tests 35 | env: 36 | PYTHONPATH: ${{ github.workspace }} 37 | run: | 38 | poetry run pytest -m "not integration" 39 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | ### Python template 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | 10 | # Distribution / packaging 11 | .Python 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | cover/ 54 | 55 | # Translations 56 | *.mo 57 | *.pot 58 | 59 | # Django stuff: 60 | *.log 61 | local_settings.py 62 | db.sqlite3 63 | db.sqlite3-journal 64 | 65 | # Flask stuff: 66 | instance/ 67 | .webassets-cache 68 | 69 | # Scrapy stuff: 70 | .scrapy 71 | 72 | # Sphinx documentation 73 | docs/_build/ 74 | 75 | # PyBuilder 76 | .pybuilder/ 77 | target/ 78 | 79 | # Jupyter Notebook 80 | .ipynb_checkpoints 81 | 82 | # IPython 83 | profile_default/ 84 | ipython_config.py 85 | 86 | # pyenv 87 | # For a library or package, you might want to ignore these files since the code is 88 | # intended to run in multiple environments; otherwise, check them in: 89 | # .python-version 90 | 91 | # pipenv 92 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 93 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 94 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 95 | # install all needed dependencies. 96 | #Pipfile.lock 97 | 98 | # poetry 99 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 100 | # This is especially recommended for binary packages to ensure reproducibility, and is more 101 | # commonly ignored for libraries. 102 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 103 | #poetry.lock 104 | 105 | # pdm 106 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 107 | #pdm.lock 108 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 109 | # in version control. 110 | # https://pdm.fming.dev/#use-with-ide 111 | .pdm.toml 112 | 113 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 114 | __pypackages__/ 115 | 116 | # Celery stuff 117 | celerybeat-schedule 118 | celerybeat.pid 119 | 120 | # SageMath parsed files 121 | *.sage.py 122 | 123 | # Environments 124 | .env 125 | .venv 126 | env/ 127 | venv/ 128 | ENV/ 129 | env.bak/ 130 | venv.bak/ 131 | 132 | # Spyder project settings 133 | .spyderproject 134 | .spyproject 135 | 136 | # Rope project settings 137 | .ropeproject 138 | 139 | # mkdocs documentation 140 | /site 141 | 142 | # mypy 143 | .mypy_cache/ 144 | .dmypy.json 145 | dmypy.json 146 | 147 | # Pyre type checker 148 | .pyre/ 149 | 150 | # pytype static type analyzer 151 | .pytype/ 152 | 153 | # Cython debug symbols 154 | cython_debug/ 155 | 156 | # PyCharm 157 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 158 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 159 | # and can be added to the global gitignore or merged into this file. For a more nuclear 160 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 161 | .idea/ 162 | .vscode/ 163 | 164 | ## Other 165 | # Cache files 166 | cache.db* 167 | 168 | # All DS_Store files 169 | .DS_Store -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Contributor Covenant Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | We as members, contributors, and leaders pledge to make participation in our 6 | community a harassment-free experience for everyone, regardless of age, body 7 | size, visible or invisible disability, ethnicity, sex characteristics, gender 8 | identity and expression, level of experience, education, socio-economic status, 9 | nationality, personal appearance, race, religion, or sexual identity 10 | and orientation. 11 | 12 | We pledge to act and interact in ways that contribute to an open, welcoming, 13 | diverse, inclusive, and healthy community. 14 | 15 | ## Our Standards 16 | 17 | Examples of behavior that contributes to a positive environment for our 18 | community include: 19 | 20 | - Demonstrating empathy and kindness toward other people 21 | - Being respectful of differing opinions, viewpoints, and experiences 22 | - Giving and gracefully accepting constructive feedback 23 | - Accepting responsibility and apologizing to those affected by our mistakes, 24 | and learning from the experience 25 | - Focusing on what is best not just for us as individuals, but for the 26 | overall community 27 | 28 | Examples of unacceptable behavior include: 29 | 30 | - The use of sexualized language or imagery, and sexual attention or 31 | advances of any kind 32 | - Trolling, insulting or derogatory comments, and personal or political attacks 33 | - Public or private harassment 34 | - Publishing others' private information, such as a physical or email 35 | address, without their explicit permission 36 | - Other conduct which could reasonably be considered inappropriate in a 37 | professional setting 38 | 39 | ## Enforcement Responsibilities 40 | 41 | Community leaders are responsible for clarifying and enforcing our standards of 42 | acceptable behavior and will take appropriate and fair corrective action in 43 | response to any behavior that they deem inappropriate, threatening, offensive, 44 | or harmful. 45 | 46 | Community leaders have the right and responsibility to remove, edit, or reject 47 | comments, commits, code, wiki edits, issues, and other contributions that are 48 | not aligned to this Code of Conduct, and will communicate reasons for moderation 49 | decisions when appropriate. 50 | 51 | ## Scope 52 | 53 | This Code of Conduct applies within all community spaces, and also applies when 54 | an individual is officially representing the community in public spaces. 55 | Examples of representing our community include using an official e-mail address, 56 | posting via an official social media account, or acting as an appointed 57 | representative at an online or offline event. 58 | 59 | ## Enforcement 60 | 61 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 62 | reported to the community leaders responsible for enforcement at 63 | founders@getzep.com. 64 | All complaints will be reviewed and investigated promptly and fairly. 65 | 66 | All community leaders are obligated to respect the privacy and security of the 67 | reporter of any incident. 68 | 69 | ## Enforcement Guidelines 70 | 71 | Community leaders will follow these Community Impact Guidelines in determining 72 | the consequences for any action they deem in violation of this Code of Conduct: 73 | 74 | ### 1. Correction 75 | 76 | **Community Impact**: Use of inappropriate language or other behavior deemed 77 | unprofessional or unwelcome in the community. 78 | 79 | **Consequence**: A private, written warning from community leaders, providing 80 | clarity around the nature of the violation and an explanation of why the 81 | behavior was inappropriate. A public apology may be requested. 82 | 83 | ### 2. Warning 84 | 85 | **Community Impact**: A violation through a single incident or series 86 | of actions. 87 | 88 | **Consequence**: A warning with consequences for continued behavior. No 89 | interaction with the people involved, including unsolicited interaction with 90 | those enforcing the Code of Conduct, for a specified period of time. This 91 | includes avoiding interactions in community spaces as well as external channels 92 | like social media. Violating these terms may lead to a temporary or 93 | permanent ban. 94 | 95 | ### 3. Temporary Ban 96 | 97 | **Community Impact**: A serious violation of community standards, including 98 | sustained inappropriate behavior. 99 | 100 | **Consequence**: A temporary ban from any sort of interaction or public 101 | communication with the community for a specified period of time. No public or 102 | private interaction with the people involved, including unsolicited interaction 103 | with those enforcing the Code of Conduct, is allowed during this period. 104 | Violating these terms may lead to a permanent ban. 105 | 106 | ### 4. Permanent Ban 107 | 108 | **Community Impact**: Demonstrating a pattern of violation of community 109 | standards, including sustained inappropriate behavior, harassment of an 110 | individual, or aggression toward or disparagement of classes of individuals. 111 | 112 | **Consequence**: A permanent ban from any sort of public interaction within 113 | the community. 114 | 115 | ## Attribution 116 | 117 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], 118 | version 2.0, available at 119 | https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. 120 | 121 | Community Impact Guidelines were inspired by [Mozilla's code of conduct 122 | enforcement ladder](https://github.com/mozilla/diversity). 123 | 124 | [homepage]: https://www.contributor-covenant.org 125 | 126 | For answers to common questions about this code of conduct, see the FAQ at 127 | https://www.contributor-covenant.org/faq. Translations are available at 128 | https://www.contributor-covenant.org/translations. 129 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to Graphiti 2 | 3 | We're thrilled you're interested in contributing to Graphiti! As firm believers in the power of open source collaboration, we're committed to building not just a tool, but a vibrant community where developers of all experience levels can make meaningful contributions. 4 | 5 | When I first joined this project, I was overwhelmed trying to figure out where to start. Someone eventually pointed me to a random "good first issue," but I later discovered there were multiple ways I could have contributed that would have better matched my skills and interests. 6 | 7 | We've restructured our contribution paths to solve this problem: 8 | 9 | # Four Ways to Get Involved 10 | 11 | ### Pick Up Existing Issues 12 | 13 | Our developers regularly tag issues with "help wanted" and "good first issue." These are pre-vetted tasks with clear scope and someone ready to help you if you get stuck. 14 | 15 | ### Create Your Own Tickets 16 | 17 | See something that needs fixing? Have an idea for an improvement? You don't need permission to identify problems. The people closest to the pain are often best positioned to describe the solution. 18 | 19 | For **feature requests**, tell us the story of what you're trying to accomplish. What are you working on? What's getting in your way? What would make your life easier? Submit these through our [GitHub issue tracker](https://github.com/getzep/graphiti/issues) with a "Feature Request" label. 20 | 21 | For **bug reports**, we need enough context to reproduce the problem. Use the [GitHub issue tracker](https://github.com/getzep/graphiti/issues) and include: 22 | 23 | - A clear title that summarizes the specific problem 24 | - What you were trying to do when you encountered the bug 25 | - What you expected to happen 26 | - What actually happened 27 | - A code sample or test case that demonstrates the issue 28 | 29 | ### Share Your Use Cases 30 | 31 | Sometimes the most valuable contribution isn't code. If you're using our project in an interesting way, add it to the [examples](https://github.com/getzep/graphiti/tree/main/examples) folder. This helps others discover new possibilities and counts as a meaningful contribution. We regularly feature compelling examples in our blog posts and videos - your work might be showcased to the broader community! 32 | 33 | ### Help Others in Discord 34 | 35 | Join our [Discord server](https://discord.com/invite/W8Kw6bsgXQ) community and pitch in at the helpdesk. Answering questions and helping troubleshoot issues is an incredibly valuable contribution that benefits everyone. The knowledge you share today saves someone hours of frustration tomorrow. 36 | 37 | ## What happens next? 38 | 39 | Once you've found an issue tagged with "good first issue" or "help wanted," or prepared an example to share, here's how to turn that into a contribution: 40 | 41 | 1. Share your approach in the issue discussion or [Discord](https://discord.com/invite/W8Kw6bsgXQ) before diving deep into code. This helps ensure your solution adheres to the architecture of Graphiti from the start and saves you from potential rework. 42 | 43 | 2. Fork the repo, make your changes in a branch, and submit a PR. We've included more detailed technical instructions below; be open to feedback during review. 44 | 45 | ## Setup 46 | 47 | 1. Fork the repository on GitHub. 48 | 2. Clone your fork locally: 49 | ``` 50 | git clone https://github.com/getzep/graphiti 51 | cd graphiti 52 | ``` 53 | 3. Set up your development environment: 54 | 55 | - Ensure you have Python 3.10+ installed. 56 | - Install Poetry: https://python-poetry.org/docs/#installation 57 | - Install project dependencies: 58 | ``` 59 | make install 60 | ``` 61 | - To run integration tests, set the appropriate environment variables 62 | 63 | ``` 64 | export TEST_OPENAI_API_KEY=... 65 | export TEST_OPENAI_MODEL=... 66 | export TEST_ANTHROPIC_API_KEY=... 67 | 68 | export NEO4J_URI=neo4j://... 69 | export NEO4J_USER=... 70 | export NEO4J_PASSWORD=... 71 | ``` 72 | 73 | ## Making Changes 74 | 75 | 1. Create a new branch for your changes: 76 | ``` 77 | git checkout -b your-branch-name 78 | ``` 79 | 2. Make your changes in the codebase. 80 | 3. Write or update tests as necessary. 81 | 4. Run the tests to ensure they pass: 82 | ``` 83 | make test 84 | ``` 85 | 5. Format your code: 86 | ``` 87 | make format 88 | ``` 89 | 6. Run linting checks: 90 | ``` 91 | make lint 92 | ``` 93 | 94 | ## Submitting Changes 95 | 96 | 1. Commit your changes: 97 | ``` 98 | git commit -m "Your detailed commit message" 99 | ``` 100 | 2. Push to your fork: 101 | ``` 102 | git push origin your-branch-name 103 | ``` 104 | 3. Submit a pull request through the GitHub website to https://github.com/getzep/graphiti. 105 | 106 | ## Pull Request Guidelines 107 | 108 | - Provide a clear title and description of your changes. 109 | - Include any relevant issue numbers in the PR description. 110 | - Ensure all tests pass and there are no linting errors. 111 | - Update documentation if you're changing functionality. 112 | 113 | ## Code Style and Quality 114 | 115 | We use several tools to maintain code quality: 116 | 117 | - Ruff for linting and formatting 118 | - Mypy for static type checking 119 | - Pytest for testing 120 | 121 | Before submitting a pull request, please run: 122 | 123 | ``` 124 | make check 125 | ``` 126 | 127 | This command will format your code, run linting checks, and execute tests. 128 | 129 | # Questions? 130 | 131 | Stuck on a contribution or have a half-formed idea? Come say hello in our [Discord server](https://discord.com/invite/W8Kw6bsgXQ). Whether you're ready to contribute or just want to learn more, we're happy to have you! It's faster than GitHub issues and you'll find both maintainers and fellow contributors ready to help. 132 | 133 | Thank you for contributing to Graphiti! 134 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Build stage 2 | FROM python:3.12-slim as builder 3 | 4 | WORKDIR /app 5 | 6 | # Install system dependencies 7 | RUN apt-get update && apt-get install -y --no-install-recommends \ 8 | gcc \ 9 | && rm -rf /var/lib/apt/lists/* 10 | 11 | # Install Poetry 12 | RUN pip install --no-cache-dir poetry 13 | 14 | # Copy only the files needed for installation 15 | COPY ./pyproject.toml ./poetry.lock* ./README.md /app/ 16 | COPY ./graphiti_core /app/graphiti_core 17 | COPY ./server/pyproject.toml ./server/poetry.lock* /app/server/ 18 | 19 | RUN poetry config virtualenvs.create false 20 | 21 | # Install the local package 22 | RUN poetry build && pip install dist/*.whl 23 | 24 | # Install server dependencies 25 | WORKDIR /app/server 26 | RUN poetry install --no-interaction --no-ansi --only main --no-root 27 | 28 | FROM python:3.12-slim 29 | 30 | # Copy only the necessary files from the builder stage 31 | COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages 32 | COPY --from=builder /usr/local/bin /usr/local/bin 33 | 34 | # Create the app directory and copy server files 35 | WORKDIR /app 36 | COPY ./server /app 37 | 38 | # Set environment variables 39 | ENV PYTHONUNBUFFERED=1 40 | ENV PORT=8000 41 | # Command to run the application 42 | 43 | CMD uvicorn graph_service.main:app --host 0.0.0.0 --port $PORT -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: install format lint test all check 2 | 3 | # Define variables 4 | PYTHON = python3 5 | POETRY = poetry 6 | PYTEST = $(POETRY) run pytest 7 | RUFF = $(POETRY) run ruff 8 | MYPY = $(POETRY) run mypy 9 | 10 | # Default target 11 | all: format lint test 12 | 13 | # Install dependencies 14 | install: 15 | $(POETRY) install --with dev 16 | 17 | # Format code 18 | format: 19 | $(RUFF) check --select I --fix 20 | $(RUFF) format 21 | 22 | # Lint code 23 | lint: 24 | $(RUFF) check 25 | $(MYPY) ./graphiti_core --show-column-numbers --show-error-codes --pretty 26 | 27 | # Run tests 28 | test: 29 | $(PYTEST) 30 | 31 | # Run format, lint, and test 32 | check: format lint test -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | # Security Policy 2 | 3 | ## Supported Versions 4 | 5 | Use this section to tell people about which versions of your project are 6 | currently being supported with security updates. 7 | 8 | | Version | Supported | 9 | |---------|--------------------| 10 | | 0.x | :white_check_mark: | 11 | 12 | 13 | ## Reporting a Vulnerability 14 | 15 | Please use GitHub's Private Vulnerability Reporting mechanism found in the Security section of this repo. 16 | -------------------------------------------------------------------------------- /Zep-CLA.md: -------------------------------------------------------------------------------- 1 | # Contributor License Agreement (CLA) 2 | 3 | In order to clarify the intellectual property license granted with Contributions from any person or entity, Zep Software, Inc. ("Zep") must have a Contributor License Agreement ("CLA") on file that has been signed by each Contributor, indicating agreement to the license terms below. This license is for your protection as a Contributor as well as the protection of Zep; it does not change your rights to use your own Contributions for any other purpose. 4 | 5 | You accept and agree to the following terms and conditions for Your present and future Contributions submitted to Zep. Except for the license granted herein to Zep and recipients of software distributed by Zep, You reserve all right, title, and interest in and to Your Contributions. 6 | 7 | ## Definitions 8 | 9 | **"You" (or "Your")** shall mean the copyright owner or legal entity authorized by the copyright owner that is making this Agreement with Zep. For legal entities, the entity making a Contribution and all other entities that control, are controlled by, or are under common control with that entity are considered to be a single Contributor. For the purposes of this definition, "control" means: 10 | 11 | i. the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or 12 | ii. ownership of fifty percent (50%) or more of the outstanding shares, or 13 | iii. beneficial ownership of such entity. 14 | 15 | **"Contribution"** shall mean any original work of authorship, including any modifications or additions to an existing work, that is intentionally submitted by You to Zep for inclusion in, or documentation of, any of the products owned or managed by Zep (the "Work"). For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to Zep or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, Zep for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by You as "Not a Contribution." 16 | 17 | ## Grant of Copyright License 18 | 19 | Subject to the terms and conditions of this Agreement, You hereby grant to Zep and to recipients of software distributed by Zep a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare derivative works of, publicly display, publicly perform, sublicense, and distribute Your Contributions and such derivative works. 20 | 21 | ## Grant of Patent License 22 | 23 | Subject to the terms and conditions of this Agreement, You hereby grant to Zep and to recipients of software distributed by Zep a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by You that are necessarily infringed by Your Contribution(s) alone or by combination of Your Contribution(s) with the Work to which such Contribution(s) was submitted. If any entity institutes patent litigation against You or any other entity (including a cross-claim or counterclaim in a lawsuit) alleging that your Contribution, or the Work to which you have contributed, constitutes direct or contributory patent infringement, then any patent licenses granted to that entity under this Agreement for that Contribution or Work shall terminate as of the date such litigation is filed. 24 | 25 | ## Representations 26 | 27 | You represent that you are legally entitled to grant the above license. If your employer(s) has rights to intellectual property that you create that includes your Contributions, you represent that you have received permission to make Contributions on behalf of that employer, that your employer has waived such rights for your Contributions to Zep, or that your employer has executed a separate Corporate CLA with Zep. 28 | 29 | You represent that each of Your Contributions is Your original creation (see section 7 for submissions on behalf of others). You represent that Your Contribution submissions include complete details of any third-party license or other restriction (including, but not limited to, related patents and trademarks) of which you are personally aware and which are associated with any part of Your Contributions. 30 | 31 | ## Support 32 | 33 | You are not expected to provide support for Your Contributions, except to the extent You desire to provide support. You may provide support for free, for a fee, or not at all. Unless required by applicable law or agreed to in writing, You provide Your Contributions on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. 34 | 35 | ## Third-Party Submissions 36 | 37 | Should You wish to submit work that is not Your original creation, You may submit it to Zep separately from any Contribution, identifying the complete details of its source and of any license or other restriction (including, but not limited to, related patents, trademarks, and license agreements) of which you are personally aware, and conspicuously marking the work as "Submitted on behalf of a third-party: [named here]". 38 | 39 | ## Notifications 40 | 41 | You agree to notify Zep of any facts or circumstances of which you become aware that would make these representations inaccurate in any respect. 42 | -------------------------------------------------------------------------------- /conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | # This code adds the project root directory to the Python path, allowing imports to work correctly when running tests. 5 | # Without this file, you might encounter ModuleNotFoundError when trying to import modules from your project, especially when running tests. 6 | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__)))) 7 | -------------------------------------------------------------------------------- /depot.json: -------------------------------------------------------------------------------- 1 | {"id":"v9jv1mlpwc"} 2 | -------------------------------------------------------------------------------- /docker-compose.test.yml: -------------------------------------------------------------------------------- 1 | version: '3.8' 2 | 3 | services: 4 | graph: 5 | image: graphiti-service:${GITHUB_SHA} 6 | ports: 7 | - "8000:8000" 8 | healthcheck: 9 | test: 10 | [ 11 | "CMD", 12 | "python", 13 | "-c", 14 | "import urllib.request; urllib.request.urlopen('http://localhost:8000/healthcheck')", 15 | ] 16 | interval: 10s 17 | timeout: 5s 18 | retries: 3 19 | depends_on: 20 | neo4j: 21 | condition: service_healthy 22 | environment: 23 | - OPENAI_API_KEY=${OPENAI_API_KEY} 24 | - NEO4J_URI=bolt://neo4j:${NEO4J_PORT} 25 | - NEO4J_USER=${NEO4J_USER} 26 | - NEO4J_PASSWORD=${NEO4J_PASSWORD} 27 | - PORT=8000 28 | 29 | neo4j: 30 | image: neo4j:5.22.0 31 | ports: 32 | - "7474:7474" 33 | - "${NEO4J_PORT}:${NEO4J_PORT}" 34 | healthcheck: 35 | test: wget "http://localhost:${NEO4J_PORT}" || exit 1 36 | interval: 1s 37 | timeout: 10s 38 | retries: 20 39 | start_period: 3s 40 | environment: 41 | - NEO4J_AUTH=${NEO4J_USER}/${NEO4J_PASSWORD} -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "3.8" 2 | 3 | services: 4 | graph: 5 | build: 6 | context: . 7 | ports: 8 | - "8000:8000" 9 | healthcheck: 10 | test: 11 | [ 12 | "CMD", 13 | "python", 14 | "-c", 15 | "import urllib.request; urllib.request.urlopen('http://localhost:8000/healthcheck')", 16 | ] 17 | interval: 10s 18 | timeout: 5s 19 | retries: 3 20 | depends_on: 21 | neo4j: 22 | condition: service_healthy 23 | environment: 24 | - OPENAI_API_KEY=${OPENAI_API_KEY} 25 | - NEO4J_URI=bolt://neo4j:${NEO4J_PORT} 26 | - NEO4J_USER=${NEO4J_USER} 27 | - NEO4J_PASSWORD=${NEO4J_PASSWORD} 28 | - PORT=8000 29 | neo4j: 30 | image: neo4j:5.26.2 31 | healthcheck: 32 | test: wget "http://localhost:${NEO4J_PORT}" || exit 1 33 | interval: 1s 34 | timeout: 10s 35 | retries: 20 36 | start_period: 3s 37 | ports: 38 | - "7474:7474" # HTTP 39 | - "${NEO4J_PORT}:${NEO4J_PORT}" # Bolt 40 | volumes: 41 | - neo4j_data:/data 42 | environment: 43 | - NEO4J_AUTH=${NEO4J_USER}/${NEO4J_PASSWORD} 44 | 45 | volumes: 46 | neo4j_data: 47 | -------------------------------------------------------------------------------- /ellipsis.yaml: -------------------------------------------------------------------------------- 1 | # See https://docs.ellipsis.dev for all available configurations. 2 | 3 | version: 1.3 4 | 5 | pr_address_comments: 6 | delivery: "new_commit" 7 | pr_review: 8 | auto_review_enabled: true # enable auto-review of PRs 9 | auto_summarize_pr: true # enable auto-summary of PRs 10 | confidence_threshold: 0.8 # Threshold for how confident Ellipsis needs to be in order to leave a comment, in range [0.0-1.0] 11 | rules: # customize behavior 12 | - "Ensure the copyright notice is present as the header of all Python files" 13 | - "Ensure code is idiomatic" 14 | - "Code should be DRY (Don't Repeat Yourself)" 15 | - "Extremely Complicated Code Needs Comments" 16 | - "Use Descriptive Variable and Constant Names" 17 | - "Follow the Single Responsibility Principle" 18 | - "Function and Method Naming Should Follow Consistent Patterns" 19 | - "There should no secrets or credentials in the code" 20 | - "Don't log sensitive data" -------------------------------------------------------------------------------- /examples/ecommerce/runner.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import json 19 | import logging 20 | import os 21 | import sys 22 | from datetime import datetime, timezone 23 | from pathlib import Path 24 | 25 | from dotenv import load_dotenv 26 | 27 | from graphiti_core import Graphiti 28 | from graphiti_core.nodes import EpisodeType 29 | from graphiti_core.utils.bulk_utils import RawEpisode 30 | from graphiti_core.utils.maintenance.graph_data_operations import clear_data 31 | 32 | load_dotenv() 33 | 34 | neo4j_uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687') 35 | neo4j_user = os.environ.get('NEO4J_USER', 'neo4j') 36 | neo4j_password = os.environ.get('NEO4J_PASSWORD', 'password') 37 | 38 | 39 | def setup_logging(): 40 | # Create a logger 41 | logger = logging.getLogger() 42 | logger.setLevel(logging.INFO) # Set the logging level to INFO 43 | 44 | # Create console handler and set level to INFO 45 | console_handler = logging.StreamHandler(sys.stdout) 46 | console_handler.setLevel(logging.INFO) 47 | 48 | # Create formatter 49 | formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s') 50 | 51 | # Add formatter to console handler 52 | console_handler.setFormatter(formatter) 53 | 54 | # Add console handler to logger 55 | logger.addHandler(console_handler) 56 | 57 | return logger 58 | 59 | 60 | shoe_conversation = [ 61 | "SalesBot: Hi, I'm Allbirds Assistant! How can I help you today?", 62 | "John: Hi, I'm looking for a new pair of shoes.", 63 | 'SalesBot: Of course! What kinde of material are you looking for?', 64 | "John: I'm looking for shoes made out of wool", 65 | """SalesBot: We have just what you are looking for, how do you like our Men's SuperLight Wool Runners 66 | - Dark Grey (Medium Grey Sole)? They use the SuperLight Foam technology.""", 67 | """John: Oh, actually I bought those 2 months ago, but unfortunately found out that I was allergic to wool. 68 | I think I will pass on those, maybe there is something with a retro look that you could suggest?""", 69 | """SalesBot: Im sorry to hear that! Would you be interested in Men's Couriers - 70 | (Blizzard Sole) model? We have them in Natural Black and Basin Blue colors""", 71 | 'John: Oh that is perfect, I LOVE the Natural Black color!. I will take those.', 72 | ] 73 | 74 | 75 | async def add_messages(client: Graphiti): 76 | for i, message in enumerate(shoe_conversation): 77 | await client.add_episode( 78 | name=f'Message {i}', 79 | episode_body=message, 80 | source=EpisodeType.message, 81 | reference_time=datetime.now(timezone.utc), 82 | source_description='Shoe conversation', 83 | ) 84 | 85 | 86 | async def main(): 87 | setup_logging() 88 | client = Graphiti(neo4j_uri, neo4j_user, neo4j_password) 89 | await clear_data(client.driver) 90 | await client.build_indices_and_constraints() 91 | await ingest_products_data(client) 92 | await add_messages(client) 93 | 94 | 95 | async def ingest_products_data(client: Graphiti): 96 | script_dir = Path(__file__).parent 97 | json_file_path = script_dir / '../data/manybirds_products.json' 98 | 99 | with open(json_file_path) as file: 100 | products = json.load(file)['products'] 101 | 102 | episodes: list[RawEpisode] = [ 103 | RawEpisode( 104 | name=f'Product {i}', 105 | content=str(product), 106 | source_description='Allbirds products', 107 | source=EpisodeType.json, 108 | reference_time=datetime.now(timezone.utc), 109 | ) 110 | for i, product in enumerate(products) 111 | ] 112 | 113 | for episode in episodes: 114 | await client.add_episode( 115 | episode.name, 116 | episode.content, 117 | episode.source_description, 118 | episode.reference_time, 119 | episode.source, 120 | ) 121 | 122 | 123 | asyncio.run(main()) 124 | -------------------------------------------------------------------------------- /examples/langgraph-agent/tinybirds-jess.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/examples/langgraph-agent/tinybirds-jess.png -------------------------------------------------------------------------------- /examples/podcast/podcast_runner.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import logging 19 | import os 20 | import sys 21 | from uuid import uuid4 22 | 23 | from dotenv import load_dotenv 24 | from pydantic import BaseModel, Field 25 | from transcript_parser import parse_podcast_messages 26 | 27 | from graphiti_core import Graphiti 28 | from graphiti_core.utils.maintenance.graph_data_operations import clear_data 29 | 30 | load_dotenv() 31 | 32 | neo4j_uri = os.environ.get('NEO4J_URI') or 'bolt://localhost:7687' 33 | neo4j_user = os.environ.get('NEO4J_USER') or 'neo4j' 34 | neo4j_password = os.environ.get('NEO4J_PASSWORD') or 'password' 35 | 36 | 37 | def setup_logging(): 38 | # Create a logger 39 | logger = logging.getLogger() 40 | logger.setLevel(logging.INFO) # Set the logging level to INFO 41 | 42 | # Create console handler and set level to INFO 43 | console_handler = logging.StreamHandler(sys.stdout) 44 | console_handler.setLevel(logging.INFO) 45 | 46 | # Create formatter 47 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 48 | 49 | # Add formatter to console handler 50 | console_handler.setFormatter(formatter) 51 | 52 | # Add console handler to logger 53 | logger.addHandler(console_handler) 54 | 55 | return logger 56 | 57 | 58 | class Person(BaseModel): 59 | """A human person, fictional or nonfictional.""" 60 | 61 | first_name: str | None = Field(..., description='First name') 62 | last_name: str | None = Field(..., description='Last name') 63 | occupation: str | None = Field(..., description="The person's work occupation") 64 | 65 | 66 | async def main(): 67 | setup_logging() 68 | client = Graphiti(neo4j_uri, neo4j_user, neo4j_password) 69 | await clear_data(client.driver) 70 | await client.build_indices_and_constraints() 71 | messages = parse_podcast_messages() 72 | group_id = str(uuid4()) 73 | 74 | for i, message in enumerate(messages[3:14]): 75 | episodes = await client.retrieve_episodes( 76 | message.actual_timestamp, 3, group_ids=['podcast'] 77 | ) 78 | episode_uuids = [episode.uuid for episode in episodes] 79 | 80 | await client.add_episode( 81 | name=f'Message {i}', 82 | episode_body=f'{message.speaker_name} ({message.role}): {message.content}', 83 | reference_time=message.actual_timestamp, 84 | source_description='Podcast Transcript', 85 | group_id=group_id, 86 | entity_types={'Person': Person}, 87 | previous_episode_uuids=episode_uuids, 88 | ) 89 | 90 | 91 | asyncio.run(main()) 92 | -------------------------------------------------------------------------------- /examples/podcast/transcript_parser.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | from datetime import datetime, timedelta, timezone 4 | 5 | from pydantic import BaseModel 6 | 7 | 8 | class Speaker(BaseModel): 9 | index: int 10 | name: str 11 | role: str 12 | 13 | 14 | class ParsedMessage(BaseModel): 15 | speaker_index: int 16 | speaker_name: str 17 | role: str 18 | relative_timestamp: str 19 | actual_timestamp: datetime 20 | content: str 21 | 22 | 23 | def parse_timestamp(timestamp: str) -> timedelta: 24 | if 'm' in timestamp: 25 | match = re.match(r'(\d+)m(?:\s*(\d+)s)?', timestamp) 26 | if match: 27 | minutes = int(match.group(1)) 28 | seconds = int(match.group(2)) if match.group(2) else 0 29 | return timedelta(minutes=minutes, seconds=seconds) 30 | elif 's' in timestamp: 31 | match = re.match(r'(\d+)s', timestamp) 32 | if match: 33 | seconds = int(match.group(1)) 34 | return timedelta(seconds=seconds) 35 | return timedelta() # Return 0 duration if parsing fails 36 | 37 | 38 | def parse_conversation_file(file_path: str, speakers: list[Speaker]) -> list[ParsedMessage]: 39 | with open(file_path) as file: 40 | content = file.read() 41 | 42 | messages = content.split('\n\n') 43 | speaker_dict = {speaker.index: speaker for speaker in speakers} 44 | 45 | parsed_messages: list[ParsedMessage] = [] 46 | 47 | # Find the last timestamp to determine podcast duration 48 | last_timestamp = timedelta() 49 | for message in reversed(messages): 50 | lines = message.strip().split('\n') 51 | if lines: 52 | first_line = lines[0] 53 | parts = first_line.split(':', 1) 54 | if len(parts) == 2: 55 | header = parts[0] 56 | header_parts = header.split() 57 | if len(header_parts) >= 2: 58 | timestamp = header_parts[1].strip('()') 59 | last_timestamp = parse_timestamp(timestamp) 60 | break 61 | 62 | # Calculate the start time 63 | now = datetime.now(timezone.utc) 64 | podcast_start_time = now - last_timestamp 65 | 66 | for message in messages: 67 | lines = message.strip().split('\n') 68 | if lines: 69 | first_line = lines[0] 70 | parts = first_line.split(':', 1) 71 | if len(parts) == 2: 72 | header, content = parts 73 | header_parts = header.split() 74 | if len(header_parts) >= 2: 75 | speaker_index = int(header_parts[0]) 76 | timestamp = header_parts[1].strip('()') 77 | 78 | if len(lines) > 1: 79 | content += '\n' + '\n'.join(lines[1:]) 80 | 81 | delta = parse_timestamp(timestamp) 82 | actual_time = podcast_start_time + delta 83 | 84 | speaker = speaker_dict.get(speaker_index) 85 | if speaker: 86 | speaker_name = speaker.name 87 | role = speaker.role 88 | else: 89 | speaker_name = f'Unknown Speaker {speaker_index}' 90 | role = 'Unknown' 91 | 92 | parsed_messages.append( 93 | ParsedMessage( 94 | speaker_index=speaker_index, 95 | speaker_name=speaker_name, 96 | role=role, 97 | relative_timestamp=timestamp, 98 | actual_timestamp=actual_time, 99 | content=content.strip(), 100 | ) 101 | ) 102 | 103 | return parsed_messages 104 | 105 | 106 | def parse_podcast_messages(): 107 | file_path = 'podcast_transcript.txt' 108 | script_dir = os.path.dirname(__file__) 109 | relative_path = os.path.join(script_dir, file_path) 110 | 111 | speakers = [ 112 | Speaker(index=0, name='Stephen DUBNER', role='Host'), 113 | Speaker(index=1, name='Tania Tetlow', role='Guest'), 114 | Speaker(index=4, name='Narrator', role='Narrator'), 115 | Speaker(index=5, name='Kamala Harris', role='Quoted'), 116 | Speaker(index=6, name='Unknown Speaker', role='Unknown'), 117 | Speaker(index=7, name='Unknown Speaker', role='Unknown'), 118 | Speaker(index=8, name='Unknown Speaker', role='Unknown'), 119 | Speaker(index=10, name='Unknown Speaker', role='Unknown'), 120 | ] 121 | 122 | parsed_conversation = parse_conversation_file(relative_path, speakers) 123 | print(f'Number of messages: {len(parsed_conversation)}') 124 | return parsed_conversation 125 | -------------------------------------------------------------------------------- /examples/quickstart/README.md: -------------------------------------------------------------------------------- 1 | # Graphiti Quickstart Example 2 | 3 | This example demonstrates the basic functionality of Graphiti, including: 4 | 5 | 1. Connecting to a Neo4j database 6 | 2. Initializing Graphiti indices and constraints 7 | 3. Adding episodes to the graph 8 | 4. Searching the graph with semantic and keyword matching 9 | 5. Exploring graph-based search with reranking using the top search result's source node UUID 10 | 6. Performing node search using predefined search recipes 11 | 12 | ## Prerequisites 13 | 14 | - Neo4j Desktop installed and running 15 | - A local DBMS created and started in Neo4j Desktop 16 | - Python 3.9+ 17 | - OpenAI API key (set as `OPENAI_API_KEY` environment variable) 18 | 19 | ## Setup Instructions 20 | 21 | 1. Install the required dependencies: 22 | 23 | ```bash 24 | pip install graphiti-core 25 | ``` 26 | 27 | 2. Set up environment variables: 28 | 29 | ```bash 30 | # Required for LLM and embedding 31 | export OPENAI_API_KEY=your_openai_api_key 32 | 33 | # Optional Neo4j connection parameters (defaults shown) 34 | export NEO4J_URI=bolt://localhost:7687 35 | export NEO4J_USER=neo4j 36 | export NEO4J_PASSWORD=password 37 | ``` 38 | 39 | 3. Run the example: 40 | 41 | ```bash 42 | python quickstart.py 43 | ``` 44 | 45 | ## What This Example Demonstrates 46 | 47 | - **Graph Initialization**: Setting up the Graphiti indices and constraints in Neo4j 48 | - **Adding Episodes**: Adding text content that will be analyzed and converted into knowledge graph nodes and edges 49 | - **Edge Search Functionality**: Performing hybrid searches that combine semantic similarity and BM25 retrieval to find relationships (edges) 50 | - **Graph-Aware Search**: Using the source node UUID from the top search result to rerank additional search results based on graph distance 51 | - **Node Search Using Recipes**: Using predefined search configurations like NODE_HYBRID_SEARCH_RRF to directly search for nodes rather than edges 52 | - **Result Processing**: Understanding the structure of search results including facts, nodes, and temporal metadata 53 | 54 | ## Next Steps 55 | 56 | After running this example, you can: 57 | 58 | 1. Modify the episode content to add your own information 59 | 2. Try different search queries to explore the knowledge extraction 60 | 3. Experiment with different center nodes for graph-distance-based reranking 61 | 4. Try other predefined search recipes from `graphiti_core.search.search_config_recipes` 62 | 5. Explore the more advanced examples in the other directories 63 | 64 | ## Understanding the Output 65 | 66 | ### Edge Search Results 67 | 68 | The edge search results include EntityEdge objects with: 69 | 70 | - UUID: Unique identifier for the edge 71 | - Fact: The extracted fact from the episode 72 | - Valid at/invalid at: Time period during which the fact was true (if available) 73 | - Source/target node UUIDs: Connections between entities in the knowledge graph 74 | 75 | ### Node Search Results 76 | 77 | The node search results include EntityNode objects with: 78 | 79 | - UUID: Unique identifier for the node 80 | - Name: The name of the entity 81 | - Content Summary: A summary of the node's content 82 | - Node Labels: The types of the node (e.g., Person, Organization) 83 | - Created At: When the node was created 84 | - Attributes: Additional properties associated with the node 85 | -------------------------------------------------------------------------------- /examples/quickstart/requirements.txt: -------------------------------------------------------------------------------- 1 | graphiti-core 2 | python-dotenv>=1.0.0 -------------------------------------------------------------------------------- /examples/wizard_of_oz/parser.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | 4 | 5 | def parse_wizard_of_oz(file_path): 6 | with open(file_path, encoding='utf-8') as file: 7 | content = file.read() 8 | 9 | # Split the content into chapters 10 | chapters = re.split(r'\n\n+Chapter [IVX]+\n', content)[ 11 | 1: 12 | ] # Skip the first split which is before Chapter I 13 | 14 | episodes = [] 15 | for i, chapter in enumerate(chapters, start=1): 16 | # Extract chapter title 17 | title_match = re.match(r'(.*?)\n\n', chapter) 18 | title = title_match.group(1) if title_match else f'Chapter {i}' 19 | 20 | # Remove the title from the chapter content 21 | chapter_content = chapter[len(title) :].strip() if title_match else chapter.strip() 22 | 23 | # Create episode dictionary 24 | episode = {'episode_number': i, 'title': title, 'content': chapter_content} 25 | episodes.append(episode) 26 | 27 | return episodes 28 | 29 | 30 | def get_wizard_of_oz_messages(): 31 | file_path = 'woo.txt' 32 | script_dir = os.path.dirname(__file__) 33 | relative_path = os.path.join(script_dir, file_path) 34 | # Use the function 35 | parsed_episodes = parse_wizard_of_oz(relative_path) 36 | return parsed_episodes 37 | -------------------------------------------------------------------------------- /examples/wizard_of_oz/runner.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import logging 19 | import os 20 | import sys 21 | from datetime import datetime, timedelta, timezone 22 | 23 | from dotenv import load_dotenv 24 | 25 | from examples.wizard_of_oz.parser import get_wizard_of_oz_messages 26 | from graphiti_core import Graphiti 27 | from graphiti_core.llm_client.anthropic_client import AnthropicClient 28 | from graphiti_core.llm_client.config import LLMConfig 29 | from graphiti_core.utils.maintenance.graph_data_operations import clear_data 30 | 31 | load_dotenv() 32 | 33 | neo4j_uri = os.environ.get('NEO4J_URI') or 'bolt://localhost:7687' 34 | neo4j_user = os.environ.get('NEO4J_USER') or 'neo4j' 35 | neo4j_password = os.environ.get('NEO4J_PASSWORD') or 'password' 36 | 37 | 38 | def setup_logging(): 39 | # Create a logger 40 | logger = logging.getLogger() 41 | logger.setLevel(logging.INFO) # Set the logging level to INFO 42 | 43 | # Create console handler and set level to INFO 44 | console_handler = logging.StreamHandler(sys.stdout) 45 | console_handler.setLevel(logging.INFO) 46 | 47 | # Create formatter 48 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 49 | 50 | # Add formatter to console handler 51 | console_handler.setFormatter(formatter) 52 | 53 | # Add console handler to logger 54 | logger.addHandler(console_handler) 55 | 56 | return logger 57 | 58 | 59 | async def main(): 60 | setup_logging() 61 | llm_client = AnthropicClient(LLMConfig(api_key=os.environ.get('ANTHROPIC_API_KEY'))) 62 | client = Graphiti(neo4j_uri, neo4j_user, neo4j_password, llm_client) 63 | messages = get_wizard_of_oz_messages() 64 | print(messages) 65 | print(len(messages)) 66 | now = datetime.now(timezone.utc) 67 | # episodes: list[BulkEpisode] = [ 68 | # BulkEpisode( 69 | # name=f'Chapter {i + 1}', 70 | # content=chapter['content'], 71 | # source_description='Wizard of Oz Transcript', 72 | # episode_type='string', 73 | # reference_time=now + timedelta(seconds=i * 10), 74 | # ) 75 | # for i, chapter in enumerate(messages[0:50]) 76 | # ] 77 | 78 | # await clear_data(client.driver) 79 | # await client.build_indices_and_constraints() 80 | # await client.add_episode_bulk(episodes) 81 | 82 | await clear_data(client.driver) 83 | await client.build_indices_and_constraints() 84 | for i, chapter in enumerate(messages): 85 | await client.add_episode( 86 | name=f'Chapter {i + 1}', 87 | episode_body=chapter['content'], 88 | source_description='Wizard of Oz Transcript', 89 | reference_time=now + timedelta(seconds=i * 10), 90 | ) 91 | 92 | 93 | asyncio.run(main()) 94 | -------------------------------------------------------------------------------- /graphiti_core/__init__.py: -------------------------------------------------------------------------------- 1 | from .graphiti import Graphiti 2 | 3 | __all__ = ['Graphiti'] 4 | -------------------------------------------------------------------------------- /graphiti_core/cross_encoder/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2025, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from .client import CrossEncoderClient 18 | from .openai_reranker_client import OpenAIRerankerClient 19 | 20 | __all__ = ['CrossEncoderClient', 'OpenAIRerankerClient'] 21 | -------------------------------------------------------------------------------- /graphiti_core/cross_encoder/bge_reranker_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | 19 | from sentence_transformers import CrossEncoder 20 | 21 | from graphiti_core.cross_encoder.client import CrossEncoderClient 22 | 23 | 24 | class BGERerankerClient(CrossEncoderClient): 25 | def __init__(self): 26 | self.model = CrossEncoder('BAAI/bge-reranker-v2-m3') 27 | 28 | async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]: 29 | if not passages: 30 | return [] 31 | 32 | input_pairs = [[query, passage] for passage in passages] 33 | 34 | # Run the synchronous predict method in an executor 35 | loop = asyncio.get_running_loop() 36 | scores = await loop.run_in_executor(None, self.model.predict, input_pairs) 37 | 38 | ranked_passages = sorted( 39 | [(passage, float(score)) for passage, score in zip(passages, scores, strict=False)], 40 | key=lambda x: x[1], 41 | reverse=True, 42 | ) 43 | 44 | return ranked_passages 45 | -------------------------------------------------------------------------------- /graphiti_core/cross_encoder/client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from abc import ABC, abstractmethod 18 | 19 | 20 | class CrossEncoderClient(ABC): 21 | """ 22 | CrossEncoderClient is an abstract base class that defines the interface 23 | for cross-encoder models used for ranking passages based on their relevance to a query. 24 | It allows for different implementations of cross-encoder models to be used interchangeably. 25 | """ 26 | 27 | @abstractmethod 28 | async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]: 29 | """ 30 | Rank the given passages based on their relevance to the query. 31 | 32 | Args: 33 | query (str): The query string. 34 | passages (list[str]): A list of passages to rank. 35 | 36 | Returns: 37 | list[tuple[str, float]]: A list of tuples containing the passage and its score, 38 | sorted in descending order of relevance. 39 | """ 40 | pass 41 | -------------------------------------------------------------------------------- /graphiti_core/cross_encoder/openai_reranker_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging 18 | from typing import Any 19 | 20 | import numpy as np 21 | import openai 22 | from openai import AsyncAzureOpenAI, AsyncOpenAI 23 | 24 | from ..helpers import semaphore_gather 25 | from ..llm_client import LLMConfig, RateLimitError 26 | from ..prompts import Message 27 | from .client import CrossEncoderClient 28 | 29 | logger = logging.getLogger(__name__) 30 | 31 | DEFAULT_MODEL = 'gpt-4.1-nano' 32 | 33 | 34 | class OpenAIRerankerClient(CrossEncoderClient): 35 | def __init__( 36 | self, 37 | config: LLMConfig | None = None, 38 | client: AsyncOpenAI | AsyncAzureOpenAI | None = None, 39 | ): 40 | """ 41 | Initialize the OpenAIRerankerClient with the provided configuration and client. 42 | 43 | This reranker uses the OpenAI API to run a simple boolean classifier prompt concurrently 44 | for each passage. Log-probabilities are used to rank the passages. 45 | 46 | Args: 47 | config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens. 48 | client (AsyncOpenAI | AsyncAzureOpenAI | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created. 49 | """ 50 | if config is None: 51 | config = LLMConfig() 52 | 53 | self.config = config 54 | if client is None: 55 | self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url) 56 | else: 57 | self.client = client 58 | 59 | async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]: 60 | openai_messages_list: Any = [ 61 | [ 62 | Message( 63 | role='system', 64 | content='You are an expert tasked with determining whether the passage is relevant to the query', 65 | ), 66 | Message( 67 | role='user', 68 | content=f""" 69 | Respond with "True" if PASSAGE is relevant to QUERY and "False" otherwise. 70 | 71 | {passage} 72 | 73 | 74 | {query} 75 | 76 | """, 77 | ), 78 | ] 79 | for passage in passages 80 | ] 81 | try: 82 | responses = await semaphore_gather( 83 | *[ 84 | self.client.chat.completions.create( 85 | model=DEFAULT_MODEL, 86 | messages=openai_messages, 87 | temperature=0, 88 | max_tokens=1, 89 | logit_bias={'6432': 1, '7983': 1}, 90 | logprobs=True, 91 | top_logprobs=2, 92 | ) 93 | for openai_messages in openai_messages_list 94 | ] 95 | ) 96 | 97 | responses_top_logprobs = [ 98 | response.choices[0].logprobs.content[0].top_logprobs 99 | if response.choices[0].logprobs is not None 100 | and response.choices[0].logprobs.content is not None 101 | else [] 102 | for response in responses 103 | ] 104 | scores: list[float] = [] 105 | for top_logprobs in responses_top_logprobs: 106 | if len(top_logprobs) == 0: 107 | continue 108 | norm_logprobs = np.exp(top_logprobs[0].logprob) 109 | if bool(top_logprobs[0].token): 110 | scores.append(norm_logprobs) 111 | else: 112 | scores.append(1 - norm_logprobs) 113 | 114 | results = [(passage, score) for passage, score in zip(passages, scores, strict=True)] 115 | results.sort(reverse=True, key=lambda x: x[1]) 116 | return results 117 | except openai.RateLimitError as e: 118 | raise RateLimitError from e 119 | except Exception as e: 120 | logger.error(f'Error in generating LLM response: {e}') 121 | raise 122 | -------------------------------------------------------------------------------- /graphiti_core/embedder/__init__.py: -------------------------------------------------------------------------------- 1 | from .client import EmbedderClient 2 | from .openai import OpenAIEmbedder, OpenAIEmbedderConfig 3 | 4 | __all__ = [ 5 | 'EmbedderClient', 6 | 'OpenAIEmbedder', 7 | 'OpenAIEmbedderConfig', 8 | ] 9 | -------------------------------------------------------------------------------- /graphiti_core/embedder/client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from abc import ABC, abstractmethod 18 | from collections.abc import Iterable 19 | 20 | from pydantic import BaseModel, Field 21 | 22 | EMBEDDING_DIM = 1024 23 | 24 | 25 | class EmbedderConfig(BaseModel): 26 | embedding_dim: int = Field(default=EMBEDDING_DIM, frozen=True) 27 | 28 | 29 | class EmbedderClient(ABC): 30 | @abstractmethod 31 | async def create( 32 | self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]] 33 | ) -> list[float]: 34 | pass 35 | 36 | async def create_batch(self, input_data_list: list[str]) -> list[list[float]]: 37 | raise NotImplementedError() 38 | -------------------------------------------------------------------------------- /graphiti_core/embedder/gemini.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Iterable 18 | 19 | from google import genai # type: ignore 20 | from google.genai import types # type: ignore 21 | from pydantic import Field 22 | 23 | from .client import EmbedderClient, EmbedderConfig 24 | 25 | DEFAULT_EMBEDDING_MODEL = 'embedding-001' 26 | 27 | 28 | class GeminiEmbedderConfig(EmbedderConfig): 29 | embedding_model: str = Field(default=DEFAULT_EMBEDDING_MODEL) 30 | api_key: str | None = None 31 | 32 | 33 | class GeminiEmbedder(EmbedderClient): 34 | """ 35 | Google Gemini Embedder Client 36 | """ 37 | 38 | def __init__(self, config: GeminiEmbedderConfig | None = None): 39 | if config is None: 40 | config = GeminiEmbedderConfig() 41 | self.config = config 42 | 43 | # Configure the Gemini API 44 | self.client = genai.Client( 45 | api_key=config.api_key, 46 | ) 47 | 48 | async def create( 49 | self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]] 50 | ) -> list[float]: 51 | """ 52 | Create embeddings for the given input data using Google's Gemini embedding model. 53 | 54 | Args: 55 | input_data: The input data to create embeddings for. Can be a string, list of strings, 56 | or an iterable of integers or iterables of integers. 57 | 58 | Returns: 59 | A list of floats representing the embedding vector. 60 | """ 61 | # Generate embeddings 62 | result = await self.client.aio.models.embed_content( 63 | model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL, 64 | contents=[input_data], # type: ignore[arg-type] # mypy fails on broad union type 65 | config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim), 66 | ) 67 | 68 | if not result.embeddings or len(result.embeddings) == 0 or not result.embeddings[0].values: 69 | raise ValueError('No embeddings returned from Gemini API in create()') 70 | 71 | return result.embeddings[0].values 72 | 73 | async def create_batch(self, input_data_list: list[str]) -> list[list[float]]: 74 | # Generate embeddings 75 | result = await self.client.aio.models.embed_content( 76 | model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL, 77 | contents=input_data_list, # type: ignore[arg-type] # mypy fails on broad union type 78 | config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim), 79 | ) 80 | 81 | if not result.embeddings or len(result.embeddings) == 0: 82 | raise Exception('No embeddings returned') 83 | 84 | embeddings = [] 85 | for embedding in result.embeddings: 86 | if not embedding.values: 87 | raise ValueError('Empty embedding values returned') 88 | embeddings.append(embedding.values) 89 | return embeddings 90 | -------------------------------------------------------------------------------- /graphiti_core/embedder/openai.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Iterable 18 | 19 | from openai import AsyncAzureOpenAI, AsyncOpenAI 20 | from openai.types import EmbeddingModel 21 | 22 | from .client import EmbedderClient, EmbedderConfig 23 | 24 | DEFAULT_EMBEDDING_MODEL = 'text-embedding-3-small' 25 | 26 | 27 | class OpenAIEmbedderConfig(EmbedderConfig): 28 | embedding_model: EmbeddingModel | str = DEFAULT_EMBEDDING_MODEL 29 | api_key: str | None = None 30 | base_url: str | None = None 31 | 32 | 33 | class OpenAIEmbedder(EmbedderClient): 34 | """ 35 | OpenAI Embedder Client 36 | 37 | This client supports both AsyncOpenAI and AsyncAzureOpenAI clients. 38 | """ 39 | 40 | def __init__( 41 | self, 42 | config: OpenAIEmbedderConfig | None = None, 43 | client: AsyncOpenAI | AsyncAzureOpenAI | None = None, 44 | ): 45 | if config is None: 46 | config = OpenAIEmbedderConfig() 47 | self.config = config 48 | 49 | if client is not None: 50 | self.client = client 51 | else: 52 | self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url) 53 | 54 | async def create( 55 | self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]] 56 | ) -> list[float]: 57 | result = await self.client.embeddings.create( 58 | input=input_data, model=self.config.embedding_model 59 | ) 60 | return result.data[0].embedding[: self.config.embedding_dim] 61 | 62 | async def create_batch(self, input_data_list: list[str]) -> list[list[float]]: 63 | result = await self.client.embeddings.create( 64 | input=input_data_list, model=self.config.embedding_model 65 | ) 66 | return [embedding.embedding[: self.config.embedding_dim] for embedding in result.data] 67 | -------------------------------------------------------------------------------- /graphiti_core/embedder/voyage.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Iterable 18 | 19 | import voyageai # type: ignore 20 | from pydantic import Field 21 | 22 | from .client import EmbedderClient, EmbedderConfig 23 | 24 | DEFAULT_EMBEDDING_MODEL = 'voyage-3' 25 | 26 | 27 | class VoyageAIEmbedderConfig(EmbedderConfig): 28 | embedding_model: str = Field(default=DEFAULT_EMBEDDING_MODEL) 29 | api_key: str | None = None 30 | 31 | 32 | class VoyageAIEmbedder(EmbedderClient): 33 | """ 34 | VoyageAI Embedder Client 35 | """ 36 | 37 | def __init__(self, config: VoyageAIEmbedderConfig | None = None): 38 | if config is None: 39 | config = VoyageAIEmbedderConfig() 40 | self.config = config 41 | self.client = voyageai.AsyncClient(api_key=config.api_key) 42 | 43 | async def create( 44 | self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]] 45 | ) -> list[float]: 46 | if isinstance(input_data, str): 47 | input_list = [input_data] 48 | elif isinstance(input_data, list): 49 | input_list = [str(i) for i in input_data if i] 50 | else: 51 | input_list = [str(i) for i in input_data if i is not None] 52 | 53 | input_list = [i for i in input_list if i] 54 | if len(input_list) == 0: 55 | return [] 56 | 57 | result = await self.client.embed(input_list, model=self.config.embedding_model) 58 | return [float(x) for x in result.embeddings[0][: self.config.embedding_dim]] 59 | 60 | async def create_batch(self, input_data_list: list[str]) -> list[list[float]]: 61 | result = await self.client.embed(input_data_list, model=self.config.embedding_model) 62 | return [ 63 | [float(x) for x in embedding[: self.config.embedding_dim]] 64 | for embedding in result.embeddings 65 | ] 66 | -------------------------------------------------------------------------------- /graphiti_core/errors.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | 18 | class GraphitiError(Exception): 19 | """Base exception class for Graphiti Core.""" 20 | 21 | 22 | class EdgeNotFoundError(GraphitiError): 23 | """Raised when an edge is not found.""" 24 | 25 | def __init__(self, uuid: str): 26 | self.message = f'edge {uuid} not found' 27 | super().__init__(self.message) 28 | 29 | 30 | class EdgesNotFoundError(GraphitiError): 31 | """Raised when a list of edges is not found.""" 32 | 33 | def __init__(self, uuids: list[str]): 34 | self.message = f'None of the edges for {uuids} were found.' 35 | super().__init__(self.message) 36 | 37 | 38 | class GroupsEdgesNotFoundError(GraphitiError): 39 | """Raised when no edges are found for a list of group ids.""" 40 | 41 | def __init__(self, group_ids: list[str]): 42 | self.message = f'no edges found for group ids {group_ids}' 43 | super().__init__(self.message) 44 | 45 | 46 | class GroupsNodesNotFoundError(GraphitiError): 47 | """Raised when no nodes are found for a list of group ids.""" 48 | 49 | def __init__(self, group_ids: list[str]): 50 | self.message = f'no nodes found for group ids {group_ids}' 51 | super().__init__(self.message) 52 | 53 | 54 | class NodeNotFoundError(GraphitiError): 55 | """Raised when a node is not found.""" 56 | 57 | def __init__(self, uuid: str): 58 | self.message = f'node {uuid} not found' 59 | super().__init__(self.message) 60 | 61 | 62 | class SearchRerankerError(GraphitiError): 63 | """Raised when a node is not found.""" 64 | 65 | def __init__(self, text: str): 66 | self.message = text 67 | super().__init__(self.message) 68 | 69 | 70 | class EntityTypeValidationError(GraphitiError): 71 | """Raised when an entity type uses protected attribute names.""" 72 | 73 | def __init__(self, entity_type: str, entity_type_attribute: str): 74 | self.message = f'{entity_type_attribute} cannot be used as an attribute for {entity_type} as it is a protected attribute name.' 75 | super().__init__(self.message) 76 | -------------------------------------------------------------------------------- /graphiti_core/graphiti_types.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from neo4j import AsyncDriver 18 | from pydantic import BaseModel, ConfigDict 19 | 20 | from graphiti_core.cross_encoder import CrossEncoderClient 21 | from graphiti_core.embedder import EmbedderClient 22 | from graphiti_core.llm_client import LLMClient 23 | 24 | 25 | class GraphitiClients(BaseModel): 26 | driver: AsyncDriver 27 | llm_client: LLMClient 28 | embedder: EmbedderClient 29 | cross_encoder: CrossEncoderClient 30 | 31 | model_config = ConfigDict(arbitrary_types_allowed=True) 32 | -------------------------------------------------------------------------------- /graphiti_core/helpers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import asyncio 18 | import os 19 | from collections.abc import Coroutine 20 | from datetime import datetime 21 | 22 | import numpy as np 23 | from dotenv import load_dotenv 24 | from neo4j import time as neo4j_time 25 | from numpy._typing import NDArray 26 | from typing_extensions import LiteralString 27 | 28 | load_dotenv() 29 | 30 | DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None) 31 | USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False)) 32 | SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20)) 33 | MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0)) 34 | DEFAULT_PAGE_LIMIT = 20 35 | 36 | RUNTIME_QUERY: LiteralString = ( 37 | 'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else '' 38 | ) 39 | 40 | 41 | def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None: 42 | return neo_date.to_native() if neo_date else None 43 | 44 | 45 | def lucene_sanitize(query: str) -> str: 46 | # Escape special characters from a query before passing into Lucene 47 | # + - && || ! ( ) { } [ ] ^ " ~ * ? : \ / 48 | escape_map = str.maketrans( 49 | { 50 | '+': r'\+', 51 | '-': r'\-', 52 | '&': r'\&', 53 | '|': r'\|', 54 | '!': r'\!', 55 | '(': r'\(', 56 | ')': r'\)', 57 | '{': r'\{', 58 | '}': r'\}', 59 | '[': r'\[', 60 | ']': r'\]', 61 | '^': r'\^', 62 | '"': r'\"', 63 | '~': r'\~', 64 | '*': r'\*', 65 | '?': r'\?', 66 | ':': r'\:', 67 | '\\': r'\\', 68 | '/': r'\/', 69 | 'O': r'\O', 70 | 'R': r'\R', 71 | 'N': r'\N', 72 | 'T': r'\T', 73 | 'A': r'\A', 74 | 'D': r'\D', 75 | } 76 | ) 77 | 78 | sanitized = query.translate(escape_map) 79 | return sanitized 80 | 81 | 82 | def normalize_l2(embedding: list[float]) -> NDArray: 83 | embedding_array = np.array(embedding) 84 | norm = np.linalg.norm(embedding_array, 2, axis=0, keepdims=True) 85 | return np.where(norm == 0, embedding_array, embedding_array / norm) 86 | 87 | 88 | # Use this instead of asyncio.gather() to bound coroutines 89 | async def semaphore_gather( 90 | *coroutines: Coroutine, 91 | max_coroutines: int = SEMAPHORE_LIMIT, 92 | ): 93 | semaphore = asyncio.Semaphore(max_coroutines) 94 | 95 | async def _wrap_coroutine(coroutine): 96 | async with semaphore: 97 | return await coroutine 98 | 99 | return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines)) 100 | -------------------------------------------------------------------------------- /graphiti_core/llm_client/__init__.py: -------------------------------------------------------------------------------- 1 | from .client import LLMClient 2 | from .config import LLMConfig 3 | from .errors import RateLimitError 4 | from .openai_client import OpenAIClient 5 | 6 | __all__ = ['LLMClient', 'OpenAIClient', 'LLMConfig', 'RateLimitError'] 7 | -------------------------------------------------------------------------------- /graphiti_core/llm_client/config.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | 19 | DEFAULT_MAX_TOKENS = 8192 20 | DEFAULT_TEMPERATURE = 0 21 | 22 | 23 | class ModelSize(Enum): 24 | small = 'small' 25 | medium = 'medium' 26 | 27 | 28 | class LLMConfig: 29 | """ 30 | Configuration class for the Language Learning Model (LLM). 31 | 32 | This class encapsulates the necessary parameters to interact with an LLM API, 33 | such as OpenAI's GPT models. It stores the API key, model name, and base URL 34 | for making requests to the LLM service. 35 | """ 36 | 37 | def __init__( 38 | self, 39 | api_key: str | None = None, 40 | model: str | None = None, 41 | base_url: str | None = None, 42 | temperature: float = DEFAULT_TEMPERATURE, 43 | max_tokens: int = DEFAULT_MAX_TOKENS, 44 | small_model: str | None = None, 45 | ): 46 | """ 47 | Initialize the LLMConfig with the provided parameters. 48 | 49 | Args: 50 | api_key (str): The authentication key for accessing the LLM API. 51 | This is required for making authorized requests. 52 | 53 | model (str, optional): The specific LLM model to use for generating responses. 54 | Defaults to "gpt-4.1-mini". 55 | 56 | base_url (str, optional): The base URL of the LLM API service. 57 | Defaults to "https://api.openai.com", which is OpenAI's standard API endpoint. 58 | This can be changed if using a different provider or a custom endpoint. 59 | 60 | small_model (str, optional): The specific LLM model to use for generating responses of simpler prompts. 61 | Defaults to "gpt-4.1-nano". 62 | """ 63 | self.base_url = base_url 64 | self.api_key = api_key 65 | self.model = model 66 | self.small_model = small_model 67 | self.temperature = temperature 68 | self.max_tokens = max_tokens 69 | -------------------------------------------------------------------------------- /graphiti_core/llm_client/errors.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | 18 | class RateLimitError(Exception): 19 | """Exception raised when the rate limit is exceeded.""" 20 | 21 | def __init__(self, message='Rate limit exceeded. Please try again later.'): 22 | self.message = message 23 | super().__init__(self.message) 24 | 25 | 26 | class RefusalError(Exception): 27 | """Exception raised when the LLM refuses to generate a response.""" 28 | 29 | def __init__(self, message: str): 30 | self.message = message 31 | super().__init__(self.message) 32 | 33 | 34 | class EmptyResponseError(Exception): 35 | """Exception raised when the LLM returns an empty response.""" 36 | 37 | def __init__(self, message: str): 38 | self.message = message 39 | super().__init__(self.message) 40 | -------------------------------------------------------------------------------- /graphiti_core/llm_client/groq_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import json 18 | import logging 19 | import typing 20 | 21 | import groq 22 | from groq import AsyncGroq 23 | from groq.types.chat import ChatCompletionMessageParam 24 | from pydantic import BaseModel 25 | 26 | from ..prompts.models import Message 27 | from .client import LLMClient 28 | from .config import LLMConfig, ModelSize 29 | from .errors import RateLimitError 30 | 31 | logger = logging.getLogger(__name__) 32 | 33 | DEFAULT_MODEL = 'llama-3.1-70b-versatile' 34 | DEFAULT_MAX_TOKENS = 2048 35 | 36 | 37 | class GroqClient(LLMClient): 38 | def __init__(self, config: LLMConfig | None = None, cache: bool = False): 39 | if config is None: 40 | config = LLMConfig(max_tokens=DEFAULT_MAX_TOKENS) 41 | elif config.max_tokens is None: 42 | config.max_tokens = DEFAULT_MAX_TOKENS 43 | super().__init__(config, cache) 44 | 45 | self.client = AsyncGroq(api_key=config.api_key) 46 | 47 | async def _generate_response( 48 | self, 49 | messages: list[Message], 50 | response_model: type[BaseModel] | None = None, 51 | max_tokens: int = DEFAULT_MAX_TOKENS, 52 | model_size: ModelSize = ModelSize.medium, 53 | ) -> dict[str, typing.Any]: 54 | msgs: list[ChatCompletionMessageParam] = [] 55 | for m in messages: 56 | if m.role == 'user': 57 | msgs.append({'role': 'user', 'content': m.content}) 58 | elif m.role == 'system': 59 | msgs.append({'role': 'system', 'content': m.content}) 60 | try: 61 | response = await self.client.chat.completions.create( 62 | model=self.model or DEFAULT_MODEL, 63 | messages=msgs, 64 | temperature=self.temperature, 65 | max_tokens=max_tokens or self.max_tokens, 66 | response_format={'type': 'json_object'}, 67 | ) 68 | result = response.choices[0].message.content or '' 69 | return json.loads(result) 70 | except groq.RateLimitError as e: 71 | raise RateLimitError from e 72 | except Exception as e: 73 | logger.error(f'Error in generating LLM response: {e}') 74 | raise 75 | -------------------------------------------------------------------------------- /graphiti_core/llm_client/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging 18 | from time import time 19 | 20 | from graphiti_core.embedder.client import EmbedderClient 21 | 22 | logger = logging.getLogger(__name__) 23 | 24 | 25 | async def generate_embedding(embedder: EmbedderClient, text: str): 26 | start = time() 27 | 28 | text = text.replace('\n', ' ') 29 | embedding = await embedder.create(input_data=[text]) 30 | 31 | end = time() 32 | logger.debug(f'embedded text of length {len(text)} in {end - start} ms') 33 | 34 | return embedding 35 | -------------------------------------------------------------------------------- /graphiti_core/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/models/__init__.py -------------------------------------------------------------------------------- /graphiti_core/models/edges/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/models/edges/__init__.py -------------------------------------------------------------------------------- /graphiti_core/models/edges/edge_db_queries.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | EPISODIC_EDGE_SAVE = """ 18 | MATCH (episode:Episodic {uuid: $episode_uuid}) 19 | MATCH (node:Entity {uuid: $entity_uuid}) 20 | MERGE (episode)-[r:MENTIONS {uuid: $uuid}]->(node) 21 | SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at} 22 | RETURN r.uuid AS uuid""" 23 | 24 | EPISODIC_EDGE_SAVE_BULK = """ 25 | UNWIND $episodic_edges AS edge 26 | MATCH (episode:Episodic {uuid: edge.source_node_uuid}) 27 | MATCH (node:Entity {uuid: edge.target_node_uuid}) 28 | MERGE (episode)-[r:MENTIONS {uuid: edge.uuid}]->(node) 29 | SET r = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at} 30 | RETURN r.uuid AS uuid 31 | """ 32 | 33 | ENTITY_EDGE_SAVE = """ 34 | MATCH (source:Entity {uuid: $source_uuid}) 35 | MATCH (target:Entity {uuid: $target_uuid}) 36 | MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target) 37 | SET r = $edge_data 38 | WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding) 39 | RETURN r.uuid AS uuid""" 40 | 41 | ENTITY_EDGE_SAVE_BULK = """ 42 | UNWIND $entity_edges AS edge 43 | MATCH (source:Entity {uuid: edge.source_node_uuid}) 44 | MATCH (target:Entity {uuid: edge.target_node_uuid}) 45 | MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target) 46 | SET r = edge 47 | WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding) 48 | RETURN edge.uuid AS uuid 49 | """ 50 | 51 | COMMUNITY_EDGE_SAVE = """ 52 | MATCH (community:Community {uuid: $community_uuid}) 53 | MATCH (node:Entity | Community {uuid: $entity_uuid}) 54 | MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node) 55 | SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at} 56 | RETURN r.uuid AS uuid""" 57 | -------------------------------------------------------------------------------- /graphiti_core/models/nodes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/models/nodes/__init__.py -------------------------------------------------------------------------------- /graphiti_core/models/nodes/node_db_queries.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | EPISODIC_NODE_SAVE = """ 18 | MERGE (n:Episodic {uuid: $uuid}) 19 | SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content, 20 | entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at} 21 | RETURN n.uuid AS uuid""" 22 | 23 | EPISODIC_NODE_SAVE_BULK = """ 24 | UNWIND $episodes AS episode 25 | MERGE (n:Episodic {uuid: episode.uuid}) 26 | SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, 27 | source: episode.source, content: episode.content, 28 | entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at} 29 | RETURN n.uuid AS uuid 30 | """ 31 | 32 | ENTITY_NODE_SAVE = """ 33 | MERGE (n:Entity {uuid: $entity_data.uuid}) 34 | SET n:$($labels) 35 | SET n = $entity_data 36 | WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding) 37 | RETURN n.uuid AS uuid""" 38 | 39 | ENTITY_NODE_SAVE_BULK = """ 40 | UNWIND $nodes AS node 41 | MERGE (n:Entity {uuid: node.uuid}) 42 | SET n:$(node.labels) 43 | SET n = node 44 | WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding) 45 | RETURN n.uuid AS uuid 46 | """ 47 | 48 | COMMUNITY_NODE_SAVE = """ 49 | MERGE (n:Community {uuid: $uuid}) 50 | SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at} 51 | WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding) 52 | RETURN n.uuid AS uuid""" 53 | -------------------------------------------------------------------------------- /graphiti_core/prompts/__init__.py: -------------------------------------------------------------------------------- 1 | from .lib import prompt_library 2 | from .models import Message 3 | 4 | __all__ = ['prompt_library', 'Message'] 5 | -------------------------------------------------------------------------------- /graphiti_core/prompts/eval.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import json 18 | from typing import Any, Protocol, TypedDict 19 | 20 | from pydantic import BaseModel, Field 21 | 22 | from .models import Message, PromptFunction, PromptVersion 23 | 24 | 25 | class QueryExpansion(BaseModel): 26 | query: str = Field(..., description='query optimized for database search') 27 | 28 | 29 | class QAResponse(BaseModel): 30 | ANSWER: str = Field(..., description='how Alice would answer the question') 31 | 32 | 33 | class EvalResponse(BaseModel): 34 | is_correct: bool = Field(..., description='boolean if the answer is correct or incorrect') 35 | reasoning: str = Field( 36 | ..., description='why you determined the response was correct or incorrect' 37 | ) 38 | 39 | 40 | class EvalAddEpisodeResults(BaseModel): 41 | candidate_is_worse: bool = Field( 42 | ..., 43 | description='boolean if the baseline extraction is higher quality than the candidate extraction.', 44 | ) 45 | reasoning: str = Field( 46 | ..., description='why you determined the response was correct or incorrect' 47 | ) 48 | 49 | 50 | class Prompt(Protocol): 51 | qa_prompt: PromptVersion 52 | eval_prompt: PromptVersion 53 | query_expansion: PromptVersion 54 | eval_add_episode_results: PromptVersion 55 | 56 | 57 | class Versions(TypedDict): 58 | qa_prompt: PromptFunction 59 | eval_prompt: PromptFunction 60 | query_expansion: PromptFunction 61 | eval_add_episode_results: PromptFunction 62 | 63 | 64 | def query_expansion(context: dict[str, Any]) -> list[Message]: 65 | sys_prompt = """You are an expert at rephrasing questions into queries used in a database retrieval system""" 66 | 67 | user_prompt = f""" 68 | Bob is asking Alice a question, are you able to rephrase the question into a simpler one about Alice in the third person 69 | that maintains the relevant context? 70 | 71 | {json.dumps(context['query'])} 72 | 73 | """ 74 | return [ 75 | Message(role='system', content=sys_prompt), 76 | Message(role='user', content=user_prompt), 77 | ] 78 | 79 | 80 | def qa_prompt(context: dict[str, Any]) -> list[Message]: 81 | sys_prompt = """You are Alice and should respond to all questions from the first person perspective of Alice""" 82 | 83 | user_prompt = f""" 84 | Your task is to briefly answer the question in the way that you think Alice would answer the question. 85 | You are given the following entity summaries and facts to help you determine the answer to your question. 86 | 87 | {json.dumps(context['entity_summaries'])} 88 | 89 | 90 | {json.dumps(context['facts'])} 91 | 92 | 93 | {context['query']} 94 | 95 | """ 96 | return [ 97 | Message(role='system', content=sys_prompt), 98 | Message(role='user', content=user_prompt), 99 | ] 100 | 101 | 102 | def eval_prompt(context: dict[str, Any]) -> list[Message]: 103 | sys_prompt = ( 104 | """You are a judge that determines if answers to questions match a gold standard answer""" 105 | ) 106 | 107 | user_prompt = f""" 108 | Given the QUESTION and the gold standard ANSWER determine if the RESPONSE to the question is correct or incorrect. 109 | Although the RESPONSE may be more verbose, mark it as correct as long as it references the same topic 110 | as the gold standard ANSWER. Also include your reasoning for the grade. 111 | 112 | {context['query']} 113 | 114 | 115 | {context['answer']} 116 | 117 | 118 | {context['response']} 119 | 120 | """ 121 | return [ 122 | Message(role='system', content=sys_prompt), 123 | Message(role='user', content=user_prompt), 124 | ] 125 | 126 | 127 | def eval_add_episode_results(context: dict[str, Any]) -> list[Message]: 128 | sys_prompt = """You are a judge that determines whether a baseline graph building result from a list of messages is better 129 | than a candidate graph building result based on the same messages.""" 130 | 131 | user_prompt = f""" 132 | Given the following PREVIOUS MESSAGES and MESSAGE, determine if the BASELINE graph data extracted from the 133 | conversation is higher quality than the CANDIDATE graph data extracted from the conversation. 134 | 135 | Return False if the BASELINE extraction is better, and True otherwise. If the CANDIDATE extraction and 136 | BASELINE extraction are nearly identical in quality, return True. Add your reasoning for your decision to the reasoning field 137 | 138 | 139 | {context['previous_messages']} 140 | 141 | 142 | {context['message']} 143 | 144 | 145 | 146 | {context['baseline']} 147 | 148 | 149 | 150 | {context['candidate']} 151 | 152 | """ 153 | return [ 154 | Message(role='system', content=sys_prompt), 155 | Message(role='user', content=user_prompt), 156 | ] 157 | 158 | 159 | versions: Versions = { 160 | 'qa_prompt': qa_prompt, 161 | 'eval_prompt': eval_prompt, 162 | 'query_expansion': query_expansion, 163 | 'eval_add_episode_results': eval_add_episode_results, 164 | } 165 | -------------------------------------------------------------------------------- /graphiti_core/prompts/extract_edge_dates.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from typing import Any, Protocol, TypedDict 18 | 19 | from pydantic import BaseModel, Field 20 | 21 | from .models import Message, PromptFunction, PromptVersion 22 | 23 | 24 | class EdgeDates(BaseModel): 25 | valid_at: str | None = Field( 26 | None, 27 | description='The date and time when the relationship described by the edge fact became true or was established. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.', 28 | ) 29 | invalid_at: str | None = Field( 30 | None, 31 | description='The date and time when the relationship described by the edge fact stopped being true or ended. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.', 32 | ) 33 | 34 | 35 | class Prompt(Protocol): 36 | v1: PromptVersion 37 | 38 | 39 | class Versions(TypedDict): 40 | v1: PromptFunction 41 | 42 | 43 | def v1(context: dict[str, Any]) -> list[Message]: 44 | return [ 45 | Message( 46 | role='system', 47 | content='You are an AI assistant that extracts datetime information for graph edges, focusing only on dates directly related to the establishment or change of the relationship described in the edge fact.', 48 | ), 49 | Message( 50 | role='user', 51 | content=f""" 52 | 53 | {context['previous_episodes']} 54 | 55 | 56 | {context['current_episode']} 57 | 58 | 59 | {context['reference_timestamp']} 60 | 61 | 62 | 63 | {context['edge_fact']} 64 | 65 | 66 | IMPORTANT: Only extract time information if it is part of the provided fact. Otherwise ignore the time mentioned. Make sure to do your best to determine the dates if only the relative time is mentioned. (eg 10 years ago, 2 mins ago) based on the provided reference timestamp 67 | If the relationship is not of spanning nature, but you are still able to determine the dates, set the valid_at only. 68 | Definitions: 69 | - valid_at: The date and time when the relationship described by the edge fact became true or was established. 70 | - invalid_at: The date and time when the relationship described by the edge fact stopped being true or ended. 71 | 72 | Task: 73 | Analyze the conversation and determine if there are dates that are part of the edge fact. Only set dates if they explicitly relate to the formation or alteration of the relationship itself. 74 | 75 | Guidelines: 76 | 1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SS.SSSSSSZ) for datetimes. 77 | 2. Use the reference timestamp as the current time when determining the valid_at and invalid_at dates. 78 | 3. If the fact is written in the present tense, use the Reference Timestamp for the valid_at date 79 | 4. If no temporal information is found that establishes or changes the relationship, leave the fields as null. 80 | 5. Do not infer dates from related events. Only use dates that are directly stated to establish or change the relationship. 81 | 6. For relative time mentions directly related to the relationship, calculate the actual datetime based on the reference timestamp. 82 | 7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date. 83 | 8. If only year is mentioned, use January 1st of that year at 00:00:00. 84 | 9. Always include the time zone offset (use Z for UTC if no specific time zone is mentioned). 85 | 10. A fact discussing that something is no longer true should have a valid_at according to when the negated fact became true. 86 | """, 87 | ), 88 | ] 89 | 90 | 91 | versions: Versions = {'v1': v1} 92 | -------------------------------------------------------------------------------- /graphiti_core/prompts/invalidate_edges.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from typing import Any, Protocol, TypedDict 18 | 19 | from pydantic import BaseModel, Field 20 | 21 | from .models import Message, PromptFunction, PromptVersion 22 | 23 | 24 | class InvalidatedEdges(BaseModel): 25 | contradicted_facts: list[int] = Field( 26 | ..., 27 | description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.', 28 | ) 29 | 30 | 31 | class Prompt(Protocol): 32 | v1: PromptVersion 33 | v2: PromptVersion 34 | 35 | 36 | class Versions(TypedDict): 37 | v1: PromptFunction 38 | v2: PromptFunction 39 | 40 | 41 | def v1(context: dict[str, Any]) -> list[Message]: 42 | return [ 43 | Message( 44 | role='system', 45 | content='You are an AI assistant that helps determine which relationships in a knowledge graph should be invalidated based solely on explicit contradictions in newer information.', 46 | ), 47 | Message( 48 | role='user', 49 | content=f""" 50 | Based on the provided existing edges and new edges with their timestamps, determine which relationships, if any, should be marked as expired due to contradictions or updates in the newer edges. 51 | Use the start and end dates of the edges to determine which edges are to be marked expired. 52 | Only mark a relationship as invalid if there is clear evidence from other edges that the relationship is no longer true. 53 | Do not invalidate relationships merely because they weren't mentioned in the episodes. You may use the current episode and previous episodes as well as the facts of each edge to understand the context of the relationships. 54 | 55 | Previous Episodes: 56 | {context['previous_episodes']} 57 | 58 | Current Episode: 59 | {context['current_episode']} 60 | 61 | Existing Edges (sorted by timestamp, newest first): 62 | {context['existing_edges']} 63 | 64 | New Edges: 65 | {context['new_edges']} 66 | 67 | Each edge is formatted as: "UUID | SOURCE_NODE - EDGE_NAME - TARGET_NODE (fact: EDGE_FACT), START_DATE (END_DATE, optional))" 68 | """, 69 | ), 70 | ] 71 | 72 | 73 | def v2(context: dict[str, Any]) -> list[Message]: 74 | return [ 75 | Message( 76 | role='system', 77 | content='You are an AI assistant that determines which facts contradict each other.', 78 | ), 79 | Message( 80 | role='user', 81 | content=f""" 82 | Based on the provided EXISTING FACTS and a NEW FACT, determine which existing facts the new fact contradicts. 83 | Return a list containing all ids of the facts that are contradicted by the NEW FACT. 84 | If there are no contradicted facts, return an empty list. 85 | 86 | 87 | {context['existing_edges']} 88 | 89 | 90 | 91 | {context['new_edge']} 92 | 93 | """, 94 | ), 95 | ] 96 | 97 | 98 | versions: Versions = {'v1': v1, 'v2': v2} 99 | -------------------------------------------------------------------------------- /graphiti_core/prompts/lib.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from typing import Any, Protocol, TypedDict 18 | 19 | from .dedupe_edges import Prompt as DedupeEdgesPrompt 20 | from .dedupe_edges import Versions as DedupeEdgesVersions 21 | from .dedupe_edges import versions as dedupe_edges_versions 22 | from .dedupe_nodes import Prompt as DedupeNodesPrompt 23 | from .dedupe_nodes import Versions as DedupeNodesVersions 24 | from .dedupe_nodes import versions as dedupe_nodes_versions 25 | from .eval import Prompt as EvalPrompt 26 | from .eval import Versions as EvalVersions 27 | from .eval import versions as eval_versions 28 | from .extract_edge_dates import Prompt as ExtractEdgeDatesPrompt 29 | from .extract_edge_dates import Versions as ExtractEdgeDatesVersions 30 | from .extract_edge_dates import versions as extract_edge_dates_versions 31 | from .extract_edges import Prompt as ExtractEdgesPrompt 32 | from .extract_edges import Versions as ExtractEdgesVersions 33 | from .extract_edges import versions as extract_edges_versions 34 | from .extract_nodes import Prompt as ExtractNodesPrompt 35 | from .extract_nodes import Versions as ExtractNodesVersions 36 | from .extract_nodes import versions as extract_nodes_versions 37 | from .invalidate_edges import Prompt as InvalidateEdgesPrompt 38 | from .invalidate_edges import Versions as InvalidateEdgesVersions 39 | from .invalidate_edges import versions as invalidate_edges_versions 40 | from .models import Message, PromptFunction 41 | from .prompt_helpers import DO_NOT_ESCAPE_UNICODE 42 | from .summarize_nodes import Prompt as SummarizeNodesPrompt 43 | from .summarize_nodes import Versions as SummarizeNodesVersions 44 | from .summarize_nodes import versions as summarize_nodes_versions 45 | 46 | 47 | class PromptLibrary(Protocol): 48 | extract_nodes: ExtractNodesPrompt 49 | dedupe_nodes: DedupeNodesPrompt 50 | extract_edges: ExtractEdgesPrompt 51 | dedupe_edges: DedupeEdgesPrompt 52 | invalidate_edges: InvalidateEdgesPrompt 53 | extract_edge_dates: ExtractEdgeDatesPrompt 54 | summarize_nodes: SummarizeNodesPrompt 55 | eval: EvalPrompt 56 | 57 | 58 | class PromptLibraryImpl(TypedDict): 59 | extract_nodes: ExtractNodesVersions 60 | dedupe_nodes: DedupeNodesVersions 61 | extract_edges: ExtractEdgesVersions 62 | dedupe_edges: DedupeEdgesVersions 63 | invalidate_edges: InvalidateEdgesVersions 64 | extract_edge_dates: ExtractEdgeDatesVersions 65 | summarize_nodes: SummarizeNodesVersions 66 | eval: EvalVersions 67 | 68 | 69 | class VersionWrapper: 70 | def __init__(self, func: PromptFunction): 71 | self.func = func 72 | 73 | def __call__(self, context: dict[str, Any]) -> list[Message]: 74 | messages = self.func(context) 75 | for message in messages: 76 | message.content += DO_NOT_ESCAPE_UNICODE if message.role == 'system' else '' 77 | return messages 78 | 79 | 80 | class PromptTypeWrapper: 81 | def __init__(self, versions: dict[str, PromptFunction]): 82 | for version, func in versions.items(): 83 | setattr(self, version, VersionWrapper(func)) 84 | 85 | 86 | class PromptLibraryWrapper: 87 | def __init__(self, library: PromptLibraryImpl): 88 | for prompt_type, versions in library.items(): 89 | setattr(self, prompt_type, PromptTypeWrapper(versions)) # type: ignore[arg-type] 90 | 91 | 92 | PROMPT_LIBRARY_IMPL: PromptLibraryImpl = { 93 | 'extract_nodes': extract_nodes_versions, 94 | 'dedupe_nodes': dedupe_nodes_versions, 95 | 'extract_edges': extract_edges_versions, 96 | 'dedupe_edges': dedupe_edges_versions, 97 | 'invalidate_edges': invalidate_edges_versions, 98 | 'extract_edge_dates': extract_edge_dates_versions, 99 | 'summarize_nodes': summarize_nodes_versions, 100 | 'eval': eval_versions, 101 | } 102 | prompt_library: PromptLibrary = PromptLibraryWrapper(PROMPT_LIBRARY_IMPL) # type: ignore[assignment] 103 | -------------------------------------------------------------------------------- /graphiti_core/prompts/models.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Callable 18 | from typing import Any, Protocol 19 | 20 | from pydantic import BaseModel 21 | 22 | 23 | class Message(BaseModel): 24 | role: str 25 | content: str 26 | 27 | 28 | class PromptVersion(Protocol): 29 | def __call__(self, context: dict[str, Any]) -> list[Message]: ... 30 | 31 | 32 | PromptFunction = Callable[[dict[str, Any]], list[Message]] 33 | -------------------------------------------------------------------------------- /graphiti_core/prompts/prompt_helpers.py: -------------------------------------------------------------------------------- 1 | DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n' 2 | -------------------------------------------------------------------------------- /graphiti_core/prompts/summarize_nodes.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import json 18 | from typing import Any, Protocol, TypedDict 19 | 20 | from pydantic import BaseModel, Field 21 | 22 | from .models import Message, PromptFunction, PromptVersion 23 | 24 | 25 | class Summary(BaseModel): 26 | summary: str = Field( 27 | ..., 28 | description='Summary containing the important information about the entity. Under 250 words', 29 | ) 30 | 31 | 32 | class SummaryDescription(BaseModel): 33 | description: str = Field(..., description='One sentence description of the provided summary') 34 | 35 | 36 | class Prompt(Protocol): 37 | summarize_pair: PromptVersion 38 | summarize_context: PromptVersion 39 | summary_description: PromptVersion 40 | 41 | 42 | class Versions(TypedDict): 43 | summarize_pair: PromptFunction 44 | summarize_context: PromptFunction 45 | summary_description: PromptFunction 46 | 47 | 48 | def summarize_pair(context: dict[str, Any]) -> list[Message]: 49 | return [ 50 | Message( 51 | role='system', 52 | content='You are a helpful assistant that combines summaries.', 53 | ), 54 | Message( 55 | role='user', 56 | content=f""" 57 | Synthesize the information from the following two summaries into a single succinct summary. 58 | 59 | Summaries must be under 250 words. 60 | 61 | Summaries: 62 | {json.dumps(context['node_summaries'], indent=2)} 63 | """, 64 | ), 65 | ] 66 | 67 | 68 | def summarize_context(context: dict[str, Any]) -> list[Message]: 69 | return [ 70 | Message( 71 | role='system', 72 | content='You are a helpful assistant that extracts entity properties from the provided text.', 73 | ), 74 | Message( 75 | role='user', 76 | content=f""" 77 | 78 | 79 | {json.dumps(context['previous_episodes'], indent=2)} 80 | {json.dumps(context['episode_content'], indent=2)} 81 | 82 | 83 | Given the above MESSAGES and the following ENTITY name, create a summary for the ENTITY. Your summary must only use 84 | information from the provided MESSAGES. Your summary should also only contain information relevant to the 85 | provided ENTITY. Summaries must be under 250 words. 86 | 87 | In addition, extract any values for the provided entity properties based on their descriptions. 88 | If the value of the entity property cannot be found in the current context, set the value of the property to the Python value None. 89 | 90 | Guidelines: 91 | 1. Do not hallucinate entity property values if they cannot be found in the current context. 92 | 2. Only use the provided messages, entity, and entity context to set attribute values. 93 | 94 | 95 | {context['node_name']} 96 | 97 | 98 | 99 | {context['node_summary']} 100 | 101 | 102 | 103 | {json.dumps(context['attributes'], indent=2)} 104 | 105 | """, 106 | ), 107 | ] 108 | 109 | 110 | def summary_description(context: dict[str, Any]) -> list[Message]: 111 | return [ 112 | Message( 113 | role='system', 114 | content='You are a helpful assistant that describes provided contents in a single sentence.', 115 | ), 116 | Message( 117 | role='user', 118 | content=f""" 119 | Create a short one sentence description of the summary that explains what kind of information is summarized. 120 | Summaries must be under 250 words. 121 | 122 | Summary: 123 | {json.dumps(context['summary'], indent=2)} 124 | """, 125 | ), 126 | ] 127 | 128 | 129 | versions: Versions = { 130 | 'summarize_pair': summarize_pair, 131 | 'summarize_context': summarize_context, 132 | 'summary_description': summary_description, 133 | } 134 | -------------------------------------------------------------------------------- /graphiti_core/py.typed: -------------------------------------------------------------------------------- 1 | # This file is intentionally left empty to indicate that the package is typed. 2 | -------------------------------------------------------------------------------- /graphiti_core/search/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/search/__init__.py -------------------------------------------------------------------------------- /graphiti_core/search/search_config.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from enum import Enum 18 | 19 | from pydantic import BaseModel, Field 20 | 21 | from graphiti_core.edges import EntityEdge 22 | from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode 23 | from graphiti_core.search.search_utils import ( 24 | DEFAULT_MIN_SCORE, 25 | DEFAULT_MMR_LAMBDA, 26 | MAX_SEARCH_DEPTH, 27 | ) 28 | 29 | DEFAULT_SEARCH_LIMIT = 10 30 | 31 | 32 | class EdgeSearchMethod(Enum): 33 | cosine_similarity = 'cosine_similarity' 34 | bm25 = 'bm25' 35 | bfs = 'breadth_first_search' 36 | 37 | 38 | class NodeSearchMethod(Enum): 39 | cosine_similarity = 'cosine_similarity' 40 | bm25 = 'bm25' 41 | bfs = 'breadth_first_search' 42 | 43 | 44 | class EpisodeSearchMethod(Enum): 45 | bm25 = 'bm25' 46 | 47 | 48 | class CommunitySearchMethod(Enum): 49 | cosine_similarity = 'cosine_similarity' 50 | bm25 = 'bm25' 51 | 52 | 53 | class EdgeReranker(Enum): 54 | rrf = 'reciprocal_rank_fusion' 55 | node_distance = 'node_distance' 56 | episode_mentions = 'episode_mentions' 57 | mmr = 'mmr' 58 | cross_encoder = 'cross_encoder' 59 | 60 | 61 | class NodeReranker(Enum): 62 | rrf = 'reciprocal_rank_fusion' 63 | node_distance = 'node_distance' 64 | episode_mentions = 'episode_mentions' 65 | mmr = 'mmr' 66 | cross_encoder = 'cross_encoder' 67 | 68 | 69 | class EpisodeReranker(Enum): 70 | rrf = 'reciprocal_rank_fusion' 71 | cross_encoder = 'cross_encoder' 72 | 73 | 74 | class CommunityReranker(Enum): 75 | rrf = 'reciprocal_rank_fusion' 76 | mmr = 'mmr' 77 | cross_encoder = 'cross_encoder' 78 | 79 | 80 | class EdgeSearchConfig(BaseModel): 81 | search_methods: list[EdgeSearchMethod] 82 | reranker: EdgeReranker = Field(default=EdgeReranker.rrf) 83 | sim_min_score: float = Field(default=DEFAULT_MIN_SCORE) 84 | mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA) 85 | bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH) 86 | 87 | 88 | class NodeSearchConfig(BaseModel): 89 | search_methods: list[NodeSearchMethod] 90 | reranker: NodeReranker = Field(default=NodeReranker.rrf) 91 | sim_min_score: float = Field(default=DEFAULT_MIN_SCORE) 92 | mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA) 93 | bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH) 94 | 95 | 96 | class EpisodeSearchConfig(BaseModel): 97 | search_methods: list[EpisodeSearchMethod] 98 | reranker: EpisodeReranker = Field(default=EpisodeReranker.rrf) 99 | sim_min_score: float = Field(default=DEFAULT_MIN_SCORE) 100 | mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA) 101 | bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH) 102 | 103 | 104 | class CommunitySearchConfig(BaseModel): 105 | search_methods: list[CommunitySearchMethod] 106 | reranker: CommunityReranker = Field(default=CommunityReranker.rrf) 107 | sim_min_score: float = Field(default=DEFAULT_MIN_SCORE) 108 | mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA) 109 | bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH) 110 | 111 | 112 | class SearchConfig(BaseModel): 113 | edge_config: EdgeSearchConfig | None = Field(default=None) 114 | node_config: NodeSearchConfig | None = Field(default=None) 115 | episode_config: EpisodeSearchConfig | None = Field(default=None) 116 | community_config: CommunitySearchConfig | None = Field(default=None) 117 | limit: int = Field(default=DEFAULT_SEARCH_LIMIT) 118 | reranker_min_score: float = Field(default=0) 119 | 120 | 121 | class SearchResults(BaseModel): 122 | edges: list[EntityEdge] 123 | nodes: list[EntityNode] 124 | episodes: list[EpisodicNode] 125 | communities: list[CommunityNode] 126 | -------------------------------------------------------------------------------- /graphiti_core/search/search_helpers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import json 18 | 19 | from graphiti_core.edges import EntityEdge 20 | from graphiti_core.search.search_config import SearchResults 21 | 22 | 23 | def format_edge_date_range(edge: EntityEdge) -> str: 24 | # return f"{datetime(edge.valid_at).strftime('%Y-%m-%d %H:%M:%S') if edge.valid_at else 'date unknown'} - {(edge.invalid_at.strftime('%Y-%m-%d %H:%M:%S') if edge.invalid_at else 'present')}" 25 | return f'{edge.valid_at if edge.valid_at else "date unknown"} - {(edge.invalid_at if edge.invalid_at else "present")}' 26 | 27 | 28 | def search_results_to_context_string(search_results: SearchResults) -> str: 29 | """Reformats a set of SearchResults into a single string to pass directly to an LLM as context""" 30 | fact_json = [ 31 | { 32 | 'fact': edge.fact, 33 | 'valid_at': str(edge.valid_at), 34 | 'invalid_at': str(edge.invalid_at or 'Present'), 35 | } 36 | for edge in search_results.edges 37 | ] 38 | entity_json = [ 39 | {'entity_name': node.name, 'summary': node.summary} for node in search_results.nodes 40 | ] 41 | episode_json = [ 42 | { 43 | 'source_description': episode.source_description, 44 | 'content': episode.content, 45 | } 46 | for episode in search_results.episodes 47 | ] 48 | community_json = [ 49 | {'community_name': community.name, 'summary': community.summary} 50 | for community in search_results.communities 51 | ] 52 | 53 | context_string = f""" 54 | FACTS and ENTITIES represent relevant context to the current conversation. 55 | COMMUNITIES represent a cluster of closely related entities. 56 | 57 | These are the most relevant facts and their valid and invalid dates. Facts are considered valid 58 | between their valid_at and invalid_at dates. Facts with an invalid_at date of "Present" are considered valid. 59 | 60 | {json.dumps(fact_json, indent=12)} 61 | 62 | 63 | {json.dumps(entity_json, indent=12)} 64 | 65 | 66 | {json.dumps(episode_json, indent=12)} 67 | 68 | 69 | {json.dumps(community_json, indent=12)} 70 | 71 | """ 72 | 73 | return context_string 74 | -------------------------------------------------------------------------------- /graphiti_core/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/utils/__init__.py -------------------------------------------------------------------------------- /graphiti_core/utils/datetime_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from datetime import datetime, timezone 18 | 19 | 20 | def utc_now() -> datetime: 21 | """Returns the current UTC datetime with timezone information.""" 22 | return datetime.now(timezone.utc) 23 | 24 | 25 | def ensure_utc(dt: datetime | None) -> datetime | None: 26 | """ 27 | Ensures a datetime is timezone-aware and in UTC. 28 | If the datetime is naive (no timezone), assumes it's in UTC. 29 | If the datetime has a different timezone, converts it to UTC. 30 | Returns None if input is None. 31 | """ 32 | if dt is None: 33 | return None 34 | 35 | if dt.tzinfo is None: 36 | # If datetime is naive, assume it's UTC 37 | return dt.replace(tzinfo=timezone.utc) 38 | elif dt.tzinfo != timezone.utc: 39 | # If datetime has a different timezone, convert to UTC 40 | return dt.astimezone(timezone.utc) 41 | 42 | return dt 43 | -------------------------------------------------------------------------------- /graphiti_core/utils/maintenance/__init__.py: -------------------------------------------------------------------------------- 1 | from .edge_operations import build_episodic_edges, extract_edges 2 | from .graph_data_operations import clear_data, retrieve_episodes 3 | from .node_operations import extract_nodes 4 | 5 | __all__ = [ 6 | 'extract_edges', 7 | 'build_episodic_edges', 8 | 'extract_nodes', 9 | 'clear_data', 10 | 'retrieve_episodes', 11 | ] 12 | -------------------------------------------------------------------------------- /graphiti_core/utils/maintenance/temporal_operations.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging 18 | from datetime import datetime 19 | from time import time 20 | 21 | from graphiti_core.edges import EntityEdge 22 | from graphiti_core.llm_client import LLMClient 23 | from graphiti_core.llm_client.config import ModelSize 24 | from graphiti_core.nodes import EpisodicNode 25 | from graphiti_core.prompts import prompt_library 26 | from graphiti_core.prompts.extract_edge_dates import EdgeDates 27 | from graphiti_core.prompts.invalidate_edges import InvalidatedEdges 28 | from graphiti_core.utils.datetime_utils import ensure_utc 29 | 30 | logger = logging.getLogger(__name__) 31 | 32 | 33 | async def extract_edge_dates( 34 | llm_client: LLMClient, 35 | edge: EntityEdge, 36 | current_episode: EpisodicNode, 37 | previous_episodes: list[EpisodicNode], 38 | ) -> tuple[datetime | None, datetime | None]: 39 | context = { 40 | 'edge_fact': edge.fact, 41 | 'current_episode': current_episode.content, 42 | 'previous_episodes': [ep.content for ep in previous_episodes], 43 | 'reference_timestamp': current_episode.valid_at.isoformat(), 44 | } 45 | llm_response = await llm_client.generate_response( 46 | prompt_library.extract_edge_dates.v1(context), response_model=EdgeDates 47 | ) 48 | 49 | valid_at = llm_response.get('valid_at') 50 | invalid_at = llm_response.get('invalid_at') 51 | 52 | valid_at_datetime = None 53 | invalid_at_datetime = None 54 | 55 | if valid_at: 56 | try: 57 | valid_at_datetime = ensure_utc(datetime.fromisoformat(valid_at.replace('Z', '+00:00'))) 58 | except ValueError as e: 59 | logger.warning(f'WARNING: Error parsing valid_at date: {e}. Input: {valid_at}') 60 | 61 | if invalid_at: 62 | try: 63 | invalid_at_datetime = ensure_utc( 64 | datetime.fromisoformat(invalid_at.replace('Z', '+00:00')) 65 | ) 66 | except ValueError as e: 67 | logger.warning(f'WARNING: Error parsing invalid_at date: {e}. Input: {invalid_at}') 68 | 69 | return valid_at_datetime, invalid_at_datetime 70 | 71 | 72 | async def get_edge_contradictions( 73 | llm_client: LLMClient, new_edge: EntityEdge, existing_edges: list[EntityEdge] 74 | ) -> list[EntityEdge]: 75 | start = time() 76 | 77 | new_edge_context = {'fact': new_edge.fact} 78 | existing_edge_context = [ 79 | {'id': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges) 80 | ] 81 | 82 | context = {'new_edge': new_edge_context, 'existing_edges': existing_edge_context} 83 | 84 | llm_response = await llm_client.generate_response( 85 | prompt_library.invalidate_edges.v2(context), 86 | response_model=InvalidatedEdges, 87 | model_size=ModelSize.small, 88 | ) 89 | 90 | contradicted_facts: list[int] = llm_response.get('contradicted_facts', []) 91 | 92 | contradicted_edges: list[EntityEdge] = [existing_edges[i] for i in contradicted_facts] 93 | 94 | end = time() 95 | logger.debug( 96 | f'Found invalidated edge candidates from {new_edge.fact}, in {(end - start) * 1000} ms' 97 | ) 98 | 99 | return contradicted_edges 100 | -------------------------------------------------------------------------------- /graphiti_core/utils/maintenance/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/graphiti_core/utils/maintenance/utils.py -------------------------------------------------------------------------------- /graphiti_core/utils/ontology_utils/entity_types_utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from pydantic import BaseModel 18 | 19 | from graphiti_core.errors import EntityTypeValidationError 20 | from graphiti_core.nodes import EntityNode 21 | 22 | 23 | def validate_entity_types( 24 | entity_types: dict[str, BaseModel] | None, 25 | ) -> bool: 26 | if entity_types is None: 27 | return True 28 | 29 | entity_node_field_names = EntityNode.model_fields.keys() 30 | 31 | for entity_type_name, entity_type_model in entity_types.items(): 32 | entity_type_field_names = entity_type_model.model_fields.keys() 33 | for entity_type_field_name in entity_type_field_names: 34 | if entity_type_field_name in entity_node_field_names: 35 | raise EntityTypeValidationError(entity_type_name, entity_type_field_name) 36 | 37 | return True 38 | -------------------------------------------------------------------------------- /images/arxiv-screenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/images/arxiv-screenshot.png -------------------------------------------------------------------------------- /images/graphiti-graph-intro.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/images/graphiti-graph-intro.gif -------------------------------------------------------------------------------- /images/graphiti-intro-slides-stock-2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/images/graphiti-intro-slides-stock-2.gif -------------------------------------------------------------------------------- /mcp_server/.env.example: -------------------------------------------------------------------------------- 1 | # Graphiti MCP Server Environment Configuration 2 | 3 | # Neo4j Database Configuration 4 | # These settings are used to connect to your Neo4j database 5 | NEO4J_URI=bolt://localhost:7687 6 | NEO4J_USER=neo4j 7 | NEO4J_PASSWORD=demodemo 8 | 9 | # OpenAI API Configuration 10 | # Required for LLM operations 11 | OPENAI_API_KEY=your_openai_api_key_here 12 | MODEL_NAME=gpt-4.1-mini 13 | 14 | # Optional: Only needed for non-standard OpenAI endpoints 15 | # OPENAI_BASE_URL=https://api.openai.com/v1 16 | 17 | # Optional: Group ID for namespacing graph data 18 | # GROUP_ID=my_project 19 | 20 | # Optional: Path configuration for Docker 21 | # PATH=/root/.local/bin:${PATH} 22 | 23 | # Optional: Memory settings for Neo4j (used in Docker Compose) 24 | # NEO4J_server_memory_heap_initial__size=512m 25 | # NEO4J_server_memory_heap_max__size=1G 26 | # NEO4J_server_memory_pagecache_size=512m 27 | 28 | # Azure OpenAI configuration 29 | # Optional: Only needed for Azure OpenAI endpoints 30 | # AZURE_OPENAI_ENDPOINT=your_azure_openai_endpoint_here 31 | # AZURE_OPENAI_API_VERSION=2025-01-01-preview 32 | # AZURE_OPENAI_DEPLOYMENT_NAME=gpt-4o-gpt-4o-mini-deployment 33 | # AZURE_OPENAI_EMBEDDING_API_VERSION=2023-05-15 34 | # AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME=text-embedding-3-large-deployment 35 | # AZURE_OPENAI_USE_MANAGED_IDENTITY=false 36 | -------------------------------------------------------------------------------- /mcp_server/.python-version: -------------------------------------------------------------------------------- 1 | 3.10 2 | -------------------------------------------------------------------------------- /mcp_server/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.11-slim 2 | 3 | WORKDIR /app 4 | 5 | # Install uv for package management 6 | RUN apt-get update && apt-get install -y curl && \ 7 | curl -LsSf https://astral.sh/uv/install.sh | sh && \ 8 | apt-get clean && rm -rf /var/lib/apt/lists/* 9 | 10 | # Add uv to PATH 11 | ENV PATH="/root/.local/bin:${PATH}" 12 | 13 | # Copy pyproject.toml and install dependencies 14 | COPY pyproject.toml . 15 | RUN uv sync 16 | 17 | # Copy application code 18 | COPY graphiti_mcp_server.py . 19 | 20 | EXPOSE 8000 21 | 22 | # Set environment variables 23 | ENV PYTHONUNBUFFERED=1 24 | 25 | # Command to run the application 26 | CMD ["uv", "run", "graphiti_mcp_server.py"] 27 | -------------------------------------------------------------------------------- /mcp_server/cursor_rules.md: -------------------------------------------------------------------------------- 1 | ## Instructions for Using Graphiti's MCP Tools for Agent Memory 2 | 3 | ### Before Starting Any Task 4 | 5 | - **Always search first:** Use the `search_nodes` tool to look for relevant preferences and procedures before beginning work. 6 | - **Search for facts too:** Use the `search_facts` tool to discover relationships and factual information that may be relevant to your task. 7 | - **Filter by entity type:** Specify `Preference`, `Procedure`, or `Requirement` in your node search to get targeted results. 8 | - **Review all matches:** Carefully examine any preferences, procedures, or facts that match your current task. 9 | 10 | ### Always Save New or Updated Information 11 | 12 | - **Capture requirements and preferences immediately:** When a user expresses a requirement or preference, use `add_episode` to store it right away. 13 | - _Best practice:_ Split very long requirements into shorter, logical chunks. 14 | - **Be explicit if something is an update to existing knowledge.** Only add what's changed or new to the graph. 15 | - **Document procedures clearly:** When you discover how a user wants things done, record it as a procedure. 16 | - **Record factual relationships:** When you learn about connections between entities, store these as facts. 17 | - **Be specific with categories:** Label preferences and procedures with clear categories for better retrieval later. 18 | 19 | ### During Your Work 20 | 21 | - **Respect discovered preferences:** Align your work with any preferences you've found. 22 | - **Follow procedures exactly:** If you find a procedure for your current task, follow it step by step. 23 | - **Apply relevant facts:** Use factual information to inform your decisions and recommendations. 24 | - **Stay consistent:** Maintain consistency with previously identified preferences, procedures, and facts. 25 | 26 | ### Best Practices 27 | 28 | - **Search before suggesting:** Always check if there's established knowledge before making recommendations. 29 | - **Combine node and fact searches:** For complex tasks, search both nodes and facts to build a complete picture. 30 | - **Use `center_node_uuid`:** When exploring related information, center your search around a specific node. 31 | - **Prioritize specific matches:** More specific information takes precedence over general information. 32 | - **Be proactive:** If you notice patterns in user behavior, consider storing them as preferences or procedures. 33 | 34 | **Remember:** The knowledge graph is your memory. Use it consistently to provide personalized assistance that respects the user's established preferences, procedures, and factual context. 35 | -------------------------------------------------------------------------------- /mcp_server/docker-compose.yml: -------------------------------------------------------------------------------- 1 | services: 2 | neo4j: 3 | image: neo4j:5.26.0 4 | ports: 5 | - "7474:7474" # HTTP 6 | - "7687:7687" # Bolt 7 | environment: 8 | - NEO4J_AUTH=neo4j/demodemo 9 | - NEO4J_server_memory_heap_initial__size=512m 10 | - NEO4J_server_memory_heap_max__size=1G 11 | - NEO4J_server_memory_pagecache_size=512m 12 | volumes: 13 | - neo4j_data:/data 14 | - neo4j_logs:/logs 15 | healthcheck: 16 | test: ["CMD", "wget", "-O", "/dev/null", "http://localhost:7474"] 17 | interval: 10s 18 | timeout: 5s 19 | retries: 5 20 | start_period: 30s 21 | 22 | graphiti-mcp: 23 | build: 24 | context: . 25 | dockerfile: Dockerfile 26 | env_file: 27 | - path: .env 28 | required: false # Makes the file optional. Default value is 'true' 29 | depends_on: 30 | neo4j: 31 | condition: service_healthy 32 | environment: 33 | - NEO4J_URI=bolt://neo4j:7687 34 | - NEO4J_USER=neo4j 35 | - NEO4J_PASSWORD=demodemo 36 | - OPENAI_API_KEY=${OPENAI_API_KEY} 37 | - MODEL_NAME=${MODEL_NAME} 38 | - PATH=/root/.local/bin:${PATH} 39 | ports: 40 | - "8000:8000" # Expose the MCP server via HTTP for SSE transport 41 | command: ["uv", "run", "graphiti_mcp_server.py", "--transport", "sse"] 42 | 43 | volumes: 44 | neo4j_data: 45 | neo4j_logs: 46 | -------------------------------------------------------------------------------- /mcp_server/mcp_config_sse_example.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "graphiti": { 4 | "transport": "sse", 5 | "url": "http://localhost:8000/sse" 6 | } 7 | } 8 | } -------------------------------------------------------------------------------- /mcp_server/mcp_config_stdio_example.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "graphiti": { 4 | "transport": "stdio", 5 | "command": "uv", 6 | "args": [ 7 | "run", 8 | "/ABSOLUTE/PATH/TO/graphiti_mcp_server.py", 9 | "--transport", 10 | "stdio" 11 | ], 12 | "env": { 13 | "NEO4J_URI": "bolt://localhost:7687", 14 | "NEO4J_USER": "neo4j", 15 | "NEO4J_PASSWORD": "demodemo", 16 | "OPENAI_API_KEY": "${OPENAI_API_KEY}", 17 | "MODEL_NAME": "gpt-4.1-mini" 18 | } 19 | } 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /mcp_server/pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "mcp-server" 3 | version = "0.1.0" 4 | description = "Graphiti MCP Server" 5 | readme = "README.md" 6 | requires-python = ">=3.10" 7 | dependencies = [ 8 | "mcp>=1.5.0", 9 | "openai>=1.68.2", 10 | "graphiti-core>=0.8.2", 11 | "azure-identity>=1.21.0", 12 | ] 13 | -------------------------------------------------------------------------------- /py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/py.typed -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "graphiti-core" 3 | description = "A temporal graph building library" 4 | version = "0.12.0pre3" 5 | authors = [ 6 | { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" }, 7 | { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" }, 8 | { "name" = "Daniel Chalef", "email" = "daniel@getzep.com" }, 9 | ] 10 | readme = "README.md" 11 | license = "Apache-2.0" 12 | requires-python = ">=3.10,<4" 13 | packages = [{ include = "graphiti_core", from = "." }] 14 | dependencies = [ 15 | "pydantic>=2.8.2", 16 | "neo4j>=5.23.0", 17 | "diskcache>=5.6.3", 18 | "openai>=1.53.0", 19 | "tenacity>=9.0.0", 20 | "numpy>=1.0.0", 21 | "python-dotenv>=1.0.1", 22 | ] 23 | 24 | [project.urls] 25 | Homepage = "https://help.getzep.com/graphiti/graphiti/overview" 26 | Repository = "https://github.com/getzep/graphiti" 27 | 28 | [project.optional-dependencies] 29 | anthropic = ["anthropic>=0.49.0"] 30 | groq = ["groq>=0.2.0"] 31 | google-genai = ["google-genai>=1.8.0"] 32 | 33 | [tool.poetry.group.dev.dependencies] 34 | mypy = ">=1.11.1" 35 | groq = ">=0.2.0" 36 | anthropic = ">=0.49.0" 37 | google-genai = ">=1.8.0" 38 | ipykernel = ">=6.29.5" 39 | jupyterlab = ">=4.2.4" 40 | diskcache-stubs = ">=5.6.3.6.20240818" 41 | langgraph = ">=0.2.15" 42 | langchain-anthropic = ">=0.2.4" 43 | langsmith = ">=0.1.108" 44 | langchain-openai = ">=0.2.6" 45 | sentence-transformers = ">=3.2.1" 46 | transformers = ">=4.45.2" 47 | voyageai = ">=0.2.3" 48 | pytest = ">=8.3.3" 49 | pytest-asyncio = ">=0.24.0" 50 | pytest-xdist = ">=3.6.1" 51 | ruff = ">=0.7.1" 52 | 53 | [build-system] 54 | requires = ["poetry-core"] 55 | build-backend = "poetry.core.masonry.api" 56 | 57 | [tool.poetry] 58 | requires-poetry = ">=2.0" 59 | 60 | [tool.pytest.ini_options] 61 | pythonpath = ["."] 62 | 63 | [tool.ruff] 64 | line-length = 100 65 | 66 | [tool.ruff.lint] 67 | select = [ 68 | # pycodestyle 69 | "E", 70 | # Pyflakes 71 | "F", 72 | # pyupgrade 73 | "UP", 74 | # flake8-bugbear 75 | "B", 76 | # flake8-simplify 77 | "SIM", 78 | # isort 79 | "I", 80 | ] 81 | ignore = ["E501"] 82 | 83 | [tool.ruff.format] 84 | quote-style = "single" 85 | indent-style = "space" 86 | docstring-code-format = true 87 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | markers = 3 | integration: marks tests as integration tests -------------------------------------------------------------------------------- /server/.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY= 2 | NEO4J_PORT=7687 3 | # Only used if not running a neo4j container in docker 4 | NEO4J_URI=bolt://localhost:7687 5 | NEO4J_USER=neo4j 6 | NEO4J_PASSWORD=password -------------------------------------------------------------------------------- /server/Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: install format lint test all check 2 | 3 | # Define variables 4 | PYTHON = python3 5 | POETRY = poetry 6 | PYTEST = $(POETRY) run pytest 7 | RUFF = $(POETRY) run ruff 8 | MYPY = $(POETRY) run mypy 9 | 10 | # Default target 11 | all: format lint test 12 | 13 | # Install dependencies 14 | install: 15 | $(POETRY) install --with dev 16 | 17 | # Format code 18 | format: 19 | $(RUFF) check --select I --fix 20 | $(RUFF) format 21 | 22 | # Lint code 23 | lint: 24 | $(RUFF) check 25 | $(MYPY) . --show-column-numbers --show-error-codes --pretty 26 | 27 | # Run tests 28 | test: 29 | $(PYTEST) 30 | 31 | # Run format, lint, and test 32 | check: format lint test -------------------------------------------------------------------------------- /server/README.md: -------------------------------------------------------------------------------- 1 | # graph-service 2 | 3 | Graph service is a fast api server implementing the [graphiti](https://github.com/getzep/graphiti) package. 4 | 5 | 6 | ## Running Instructions 7 | 8 | 1. Ensure you have Docker and Docker Compose installed on your system. 9 | 10 | 2. Add `zepai/graphiti:latest` to your service setup 11 | 12 | 3. Make sure to pass the following environment variables to the service 13 | 14 | ``` 15 | OPENAI_API_KEY=your_openai_api_key 16 | NEO4J_USER=your_neo4j_user 17 | NEO4J_PASSWORD=your_neo4j_password 18 | NEO4J_PORT=your_neo4j_port 19 | ``` 20 | 21 | 4. This service depends on having access to a neo4j instance, you may wish to add a neo4j image to your service setup as well. Or you may wish to use neo4j cloud or a desktop version if running this locally. 22 | 23 | An example of docker compose setup may look like this: 24 | 25 | ```yml 26 | version: '3.8' 27 | 28 | services: 29 | graph: 30 | image: zepai/graphiti:latest 31 | ports: 32 | - "8000:8000" 33 | 34 | environment: 35 | - OPENAI_API_KEY=${OPENAI_API_KEY} 36 | - NEO4J_URI=bolt://neo4j:${NEO4J_PORT} 37 | - NEO4J_USER=${NEO4J_USER} 38 | - NEO4J_PASSWORD=${NEO4J_PASSWORD} 39 | neo4j: 40 | image: neo4j:5.22.0 41 | 42 | ports: 43 | - "7474:7474" # HTTP 44 | - "${NEO4J_PORT}:${NEO4J_PORT}" # Bolt 45 | volumes: 46 | - neo4j_data:/data 47 | environment: 48 | - NEO4J_AUTH=${NEO4J_USER}/${NEO4J_PASSWORD} 49 | 50 | volumes: 51 | neo4j_data: 52 | ``` 53 | 54 | 5. Once you start the service, it will be available at `http://localhost:8000` (or the port you have specified in the docker compose file). 55 | 56 | 6. You may access the swagger docs at `http://localhost:8000/docs`. You may also access redocs at `http://localhost:8000/redoc`. 57 | 58 | 7. You may also access the neo4j browser at `http://localhost:7474` (the port depends on the neo4j instance you are using). -------------------------------------------------------------------------------- /server/graph_service/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/server/graph_service/__init__.py -------------------------------------------------------------------------------- /server/graph_service/config.py: -------------------------------------------------------------------------------- 1 | from functools import lru_cache 2 | from typing import Annotated 3 | 4 | from fastapi import Depends 5 | from pydantic import Field 6 | from pydantic_settings import BaseSettings, SettingsConfigDict # type: ignore 7 | 8 | 9 | class Settings(BaseSettings): 10 | openai_api_key: str 11 | openai_base_url: str | None = Field(None) 12 | model_name: str | None = Field(None) 13 | embedding_model_name: str | None = Field(None) 14 | neo4j_uri: str 15 | neo4j_user: str 16 | neo4j_password: str 17 | 18 | model_config = SettingsConfigDict(env_file='.env', extra='ignore') 19 | 20 | 21 | @lru_cache 22 | def get_settings(): 23 | return Settings() 24 | 25 | 26 | ZepEnvDep = Annotated[Settings, Depends(get_settings)] 27 | -------------------------------------------------------------------------------- /server/graph_service/dto/__init__.py: -------------------------------------------------------------------------------- 1 | from .common import Message, Result 2 | from .ingest import AddEntityNodeRequest, AddMessagesRequest 3 | from .retrieve import FactResult, GetMemoryRequest, GetMemoryResponse, SearchQuery, SearchResults 4 | 5 | __all__ = [ 6 | 'SearchQuery', 7 | 'Message', 8 | 'AddMessagesRequest', 9 | 'AddEntityNodeRequest', 10 | 'SearchResults', 11 | 'FactResult', 12 | 'Result', 13 | 'GetMemoryRequest', 14 | 'GetMemoryResponse', 15 | ] 16 | -------------------------------------------------------------------------------- /server/graph_service/dto/common.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | from typing import Literal 3 | 4 | from graphiti_core.utils.datetime_utils import utc_now 5 | from pydantic import BaseModel, Field 6 | 7 | 8 | class Result(BaseModel): 9 | message: str 10 | success: bool 11 | 12 | 13 | class Message(BaseModel): 14 | content: str = Field(..., description='The content of the message') 15 | uuid: str | None = Field(default=None, description='The uuid of the message (optional)') 16 | name: str = Field( 17 | default='', description='The name of the episodic node for the message (optional)' 18 | ) 19 | role_type: Literal['user', 'assistant', 'system'] = Field( 20 | ..., description='The role type of the message (user, assistant or system)' 21 | ) 22 | role: str | None = Field( 23 | description='The custom role of the message to be used alongside role_type (user name, bot name, etc.)', 24 | ) 25 | timestamp: datetime = Field(default_factory=utc_now, description='The timestamp of the message') 26 | source_description: str = Field( 27 | default='', description='The description of the source of the message' 28 | ) 29 | -------------------------------------------------------------------------------- /server/graph_service/dto/ingest.py: -------------------------------------------------------------------------------- 1 | from pydantic import BaseModel, Field 2 | 3 | from graph_service.dto.common import Message 4 | 5 | 6 | class AddMessagesRequest(BaseModel): 7 | group_id: str = Field(..., description='The group id of the messages to add') 8 | messages: list[Message] = Field(..., description='The messages to add') 9 | 10 | 11 | class AddEntityNodeRequest(BaseModel): 12 | uuid: str = Field(..., description='The uuid of the node to add') 13 | group_id: str = Field(..., description='The group id of the node to add') 14 | name: str = Field(..., description='The name of the node to add') 15 | summary: str = Field(default='', description='The summary of the node to add') 16 | -------------------------------------------------------------------------------- /server/graph_service/dto/retrieve.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime, timezone 2 | 3 | from pydantic import BaseModel, Field 4 | 5 | from graph_service.dto.common import Message 6 | 7 | 8 | class SearchQuery(BaseModel): 9 | group_ids: list[str] | None = Field( 10 | None, description='The group ids for the memories to search' 11 | ) 12 | query: str 13 | max_facts: int = Field(default=10, description='The maximum number of facts to retrieve') 14 | 15 | 16 | class FactResult(BaseModel): 17 | uuid: str 18 | name: str 19 | fact: str 20 | valid_at: datetime | None 21 | invalid_at: datetime | None 22 | created_at: datetime 23 | expired_at: datetime | None 24 | 25 | class Config: 26 | json_encoders = {datetime: lambda v: v.astimezone(timezone.utc).isoformat()} 27 | 28 | 29 | class SearchResults(BaseModel): 30 | facts: list[FactResult] 31 | 32 | 33 | class GetMemoryRequest(BaseModel): 34 | group_id: str = Field(..., description='The group id of the memory to get') 35 | max_facts: int = Field(default=10, description='The maximum number of facts to retrieve') 36 | center_node_uuid: str | None = Field( 37 | ..., description='The uuid of the node to center the retrieval on' 38 | ) 39 | messages: list[Message] = Field( 40 | ..., description='The messages to build the retrieval query from ' 41 | ) 42 | 43 | 44 | class GetMemoryResponse(BaseModel): 45 | facts: list[FactResult] = Field(..., description='The facts that were retrieved from the graph') 46 | -------------------------------------------------------------------------------- /server/graph_service/main.py: -------------------------------------------------------------------------------- 1 | from contextlib import asynccontextmanager 2 | 3 | from fastapi import FastAPI 4 | from fastapi.responses import JSONResponse 5 | 6 | from graph_service.config import get_settings 7 | from graph_service.routers import ingest, retrieve 8 | from graph_service.zep_graphiti import initialize_graphiti 9 | 10 | 11 | @asynccontextmanager 12 | async def lifespan(_: FastAPI): 13 | settings = get_settings() 14 | await initialize_graphiti(settings) 15 | yield 16 | # Shutdown 17 | # No need to close Graphiti here, as it's handled per-request 18 | 19 | 20 | app = FastAPI(lifespan=lifespan) 21 | 22 | 23 | app.include_router(retrieve.router) 24 | app.include_router(ingest.router) 25 | 26 | 27 | @app.get('/healthcheck') 28 | async def healthcheck(): 29 | return JSONResponse(content={'status': 'healthy'}, status_code=200) 30 | -------------------------------------------------------------------------------- /server/graph_service/routers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/getzep/graphiti/1ac9589cbf703bc502e58abe7884f0f16dc0a4be/server/graph_service/routers/__init__.py -------------------------------------------------------------------------------- /server/graph_service/routers/ingest.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from contextlib import asynccontextmanager 3 | from functools import partial 4 | 5 | from fastapi import APIRouter, FastAPI, status 6 | from graphiti_core.nodes import EpisodeType # type: ignore 7 | from graphiti_core.utils.maintenance.graph_data_operations import clear_data # type: ignore 8 | 9 | from graph_service.dto import AddEntityNodeRequest, AddMessagesRequest, Message, Result 10 | from graph_service.zep_graphiti import ZepGraphitiDep 11 | 12 | 13 | class AsyncWorker: 14 | def __init__(self): 15 | self.queue = asyncio.Queue() 16 | self.task = None 17 | 18 | async def worker(self): 19 | while True: 20 | try: 21 | print(f'Got a job: (size of remaining queue: {self.queue.qsize()})') 22 | job = await self.queue.get() 23 | await job() 24 | except asyncio.CancelledError: 25 | break 26 | 27 | async def start(self): 28 | self.task = asyncio.create_task(self.worker()) 29 | 30 | async def stop(self): 31 | if self.task: 32 | self.task.cancel() 33 | await self.task 34 | while not self.queue.empty(): 35 | self.queue.get_nowait() 36 | 37 | 38 | async_worker = AsyncWorker() 39 | 40 | 41 | @asynccontextmanager 42 | async def lifespan(_: FastAPI): 43 | await async_worker.start() 44 | yield 45 | await async_worker.stop() 46 | 47 | 48 | router = APIRouter(lifespan=lifespan) 49 | 50 | 51 | @router.post('/messages', status_code=status.HTTP_202_ACCEPTED) 52 | async def add_messages( 53 | request: AddMessagesRequest, 54 | graphiti: ZepGraphitiDep, 55 | ): 56 | async def add_messages_task(m: Message): 57 | await graphiti.add_episode( 58 | uuid=m.uuid, 59 | group_id=request.group_id, 60 | name=m.name, 61 | episode_body=f'{m.role or ""}({m.role_type}): {m.content}', 62 | reference_time=m.timestamp, 63 | source=EpisodeType.message, 64 | source_description=m.source_description, 65 | ) 66 | 67 | for m in request.messages: 68 | await async_worker.queue.put(partial(add_messages_task, m)) 69 | 70 | return Result(message='Messages added to processing queue', success=True) 71 | 72 | 73 | @router.post('/entity-node', status_code=status.HTTP_201_CREATED) 74 | async def add_entity_node( 75 | request: AddEntityNodeRequest, 76 | graphiti: ZepGraphitiDep, 77 | ): 78 | node = await graphiti.save_entity_node( 79 | uuid=request.uuid, 80 | group_id=request.group_id, 81 | name=request.name, 82 | summary=request.summary, 83 | ) 84 | return node 85 | 86 | 87 | @router.delete('/entity-edge/{uuid}', status_code=status.HTTP_200_OK) 88 | async def delete_entity_edge(uuid: str, graphiti: ZepGraphitiDep): 89 | await graphiti.delete_entity_edge(uuid) 90 | return Result(message='Entity Edge deleted', success=True) 91 | 92 | 93 | @router.delete('/group/{group_id}', status_code=status.HTTP_200_OK) 94 | async def delete_group(group_id: str, graphiti: ZepGraphitiDep): 95 | await graphiti.delete_group(group_id) 96 | return Result(message='Group deleted', success=True) 97 | 98 | 99 | @router.delete('/episode/{uuid}', status_code=status.HTTP_200_OK) 100 | async def delete_episode(uuid: str, graphiti: ZepGraphitiDep): 101 | await graphiti.delete_episodic_node(uuid) 102 | return Result(message='Episode deleted', success=True) 103 | 104 | 105 | @router.post('/clear', status_code=status.HTTP_200_OK) 106 | async def clear( 107 | graphiti: ZepGraphitiDep, 108 | ): 109 | await clear_data(graphiti.driver) 110 | await graphiti.build_indices_and_constraints() 111 | return Result(message='Graph cleared', success=True) 112 | -------------------------------------------------------------------------------- /server/graph_service/routers/retrieve.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime, timezone 2 | 3 | from fastapi import APIRouter, status 4 | 5 | from graph_service.dto import ( 6 | GetMemoryRequest, 7 | GetMemoryResponse, 8 | Message, 9 | SearchQuery, 10 | SearchResults, 11 | ) 12 | from graph_service.zep_graphiti import ZepGraphitiDep, get_fact_result_from_edge 13 | 14 | router = APIRouter() 15 | 16 | 17 | @router.post('/search', status_code=status.HTTP_200_OK) 18 | async def search(query: SearchQuery, graphiti: ZepGraphitiDep): 19 | relevant_edges = await graphiti.search( 20 | group_ids=query.group_ids, 21 | query=query.query, 22 | num_results=query.max_facts, 23 | ) 24 | facts = [get_fact_result_from_edge(edge) for edge in relevant_edges] 25 | return SearchResults( 26 | facts=facts, 27 | ) 28 | 29 | 30 | @router.get('/entity-edge/{uuid}', status_code=status.HTTP_200_OK) 31 | async def get_entity_edge(uuid: str, graphiti: ZepGraphitiDep): 32 | entity_edge = await graphiti.get_entity_edge(uuid) 33 | return get_fact_result_from_edge(entity_edge) 34 | 35 | 36 | @router.get('/episodes/{group_id}', status_code=status.HTTP_200_OK) 37 | async def get_episodes(group_id: str, last_n: int, graphiti: ZepGraphitiDep): 38 | episodes = await graphiti.retrieve_episodes( 39 | group_ids=[group_id], last_n=last_n, reference_time=datetime.now(timezone.utc) 40 | ) 41 | return episodes 42 | 43 | 44 | @router.post('/get-memory', status_code=status.HTTP_200_OK) 45 | async def get_memory( 46 | request: GetMemoryRequest, 47 | graphiti: ZepGraphitiDep, 48 | ): 49 | combined_query = compose_query_from_messages(request.messages) 50 | result = await graphiti.search( 51 | group_ids=[request.group_id], 52 | query=combined_query, 53 | num_results=request.max_facts, 54 | ) 55 | facts = [get_fact_result_from_edge(edge) for edge in result] 56 | return GetMemoryResponse(facts=facts) 57 | 58 | 59 | def compose_query_from_messages(messages: list[Message]): 60 | combined_query = '' 61 | for message in messages: 62 | combined_query += f'{message.role_type or ""}({message.role or ""}): {message.content}\n' 63 | return combined_query 64 | -------------------------------------------------------------------------------- /server/graph_service/zep_graphiti.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import Annotated 3 | 4 | from fastapi import Depends, HTTPException 5 | from graphiti_core import Graphiti # type: ignore 6 | from graphiti_core.edges import EntityEdge # type: ignore 7 | from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError, NodeNotFoundError 8 | from graphiti_core.llm_client import LLMClient # type: ignore 9 | from graphiti_core.nodes import EntityNode, EpisodicNode # type: ignore 10 | 11 | from graph_service.config import ZepEnvDep 12 | from graph_service.dto import FactResult 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | 17 | class ZepGraphiti(Graphiti): 18 | def __init__(self, uri: str, user: str, password: str, llm_client: LLMClient | None = None): 19 | super().__init__(uri, user, password, llm_client) 20 | 21 | async def save_entity_node(self, name: str, uuid: str, group_id: str, summary: str = ''): 22 | new_node = EntityNode( 23 | name=name, 24 | uuid=uuid, 25 | group_id=group_id, 26 | summary=summary, 27 | ) 28 | await new_node.generate_name_embedding(self.embedder) 29 | await new_node.save(self.driver) 30 | return new_node 31 | 32 | async def get_entity_edge(self, uuid: str): 33 | try: 34 | edge = await EntityEdge.get_by_uuid(self.driver, uuid) 35 | return edge 36 | except EdgeNotFoundError as e: 37 | raise HTTPException(status_code=404, detail=e.message) from e 38 | 39 | async def delete_group(self, group_id: str): 40 | try: 41 | edges = await EntityEdge.get_by_group_ids(self.driver, [group_id]) 42 | except GroupsEdgesNotFoundError: 43 | logger.warning(f'No edges found for group {group_id}') 44 | edges = [] 45 | 46 | nodes = await EntityNode.get_by_group_ids(self.driver, [group_id]) 47 | 48 | episodes = await EpisodicNode.get_by_group_ids(self.driver, [group_id]) 49 | 50 | for edge in edges: 51 | await edge.delete(self.driver) 52 | 53 | for node in nodes: 54 | await node.delete(self.driver) 55 | 56 | for episode in episodes: 57 | await episode.delete(self.driver) 58 | 59 | async def delete_entity_edge(self, uuid: str): 60 | try: 61 | edge = await EntityEdge.get_by_uuid(self.driver, uuid) 62 | await edge.delete(self.driver) 63 | except EdgeNotFoundError as e: 64 | raise HTTPException(status_code=404, detail=e.message) from e 65 | 66 | async def delete_episodic_node(self, uuid: str): 67 | try: 68 | episode = await EpisodicNode.get_by_uuid(self.driver, uuid) 69 | await episode.delete(self.driver) 70 | except NodeNotFoundError as e: 71 | raise HTTPException(status_code=404, detail=e.message) from e 72 | 73 | 74 | async def get_graphiti(settings: ZepEnvDep): 75 | client = ZepGraphiti( 76 | uri=settings.neo4j_uri, 77 | user=settings.neo4j_user, 78 | password=settings.neo4j_password, 79 | ) 80 | if settings.openai_base_url is not None: 81 | client.llm_client.config.base_url = settings.openai_base_url 82 | if settings.openai_api_key is not None: 83 | client.llm_client.config.api_key = settings.openai_api_key 84 | if settings.model_name is not None: 85 | client.llm_client.model = settings.model_name 86 | 87 | try: 88 | yield client 89 | finally: 90 | await client.close() 91 | 92 | 93 | async def initialize_graphiti(settings: ZepEnvDep): 94 | client = ZepGraphiti( 95 | uri=settings.neo4j_uri, 96 | user=settings.neo4j_user, 97 | password=settings.neo4j_password, 98 | ) 99 | await client.build_indices_and_constraints() 100 | 101 | 102 | def get_fact_result_from_edge(edge: EntityEdge): 103 | return FactResult( 104 | uuid=edge.uuid, 105 | name=edge.name, 106 | fact=edge.fact, 107 | valid_at=edge.valid_at, 108 | invalid_at=edge.invalid_at, 109 | created_at=edge.created_at, 110 | expired_at=edge.expired_at, 111 | ) 112 | 113 | 114 | ZepGraphitiDep = Annotated[ZepGraphiti, Depends(get_graphiti)] 115 | -------------------------------------------------------------------------------- /server/pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "graph-service" 3 | version = "0.1.0" 4 | description = "Zep Graph service implementing Graphiti package" 5 | authors = ["Paul Paliychuk "] 6 | readme = "README.md" 7 | packages = [{ include = "graph_service" }] 8 | 9 | 10 | [tool.poetry.dependencies] 11 | python = "^3.10" 12 | fastapi = "^0.115.0" 13 | graphiti-core = { path = "../" } 14 | pydantic-settings = "^2.4.0" 15 | uvicorn = "^0.30.6" 16 | httpx = "^0.28.1" 17 | 18 | [tool.poetry.group.dev.dependencies] 19 | pydantic = "^2.8.2" 20 | mypy = "^1.11.1" 21 | pytest = "^8.3.2" 22 | python-dotenv = "^1.0.1" 23 | pytest-asyncio = "^0.24.0" 24 | pytest-xdist = "^3.6.1" 25 | ruff = "^0.6.2" 26 | fastapi-cli = "^0.0.5" 27 | 28 | [build-system] 29 | requires = ["poetry-core"] 30 | build-backend = "poetry.core.masonry.api" 31 | 32 | [tool.pytest.ini_options] 33 | pythonpath = ["."] 34 | 35 | [tool.ruff] 36 | line-length = 100 37 | 38 | [tool.ruff.lint] 39 | select = [ 40 | # pycodestyle 41 | "E", 42 | # Pyflakes 43 | "F", 44 | # pyupgrade 45 | "UP", 46 | # flake8-bugbear 47 | "B", 48 | # flake8-simplify 49 | "SIM", 50 | # isort 51 | "I", 52 | ] 53 | ignore = ["E501"] 54 | 55 | [tool.ruff.format] 56 | quote-style = "single" 57 | indent-style = "space" 58 | docstring-code-format = true 59 | -------------------------------------------------------------------------------- /signatures/version1/cla.json: -------------------------------------------------------------------------------- 1 | { 2 | "signedContributors": [ 3 | { 4 | "name": "colombod", 5 | "id": 375556, 6 | "comment_id": 2761979440, 7 | "created_at": "2025-03-28T17:21:29Z", 8 | "repoId": 840056306, 9 | "pullRequestNo": 310 10 | }, 11 | { 12 | "name": "evanmschultz", 13 | "id": 3806601, 14 | "comment_id": 2813673237, 15 | "created_at": "2025-04-17T17:56:24Z", 16 | "repoId": 840056306, 17 | "pullRequestNo": 372 18 | }, 19 | { 20 | "name": "soichisumi", 21 | "id": 30210641, 22 | "comment_id": 2818469528, 23 | "created_at": "2025-04-21T14:02:11Z", 24 | "repoId": 840056306, 25 | "pullRequestNo": 382 26 | }, 27 | { 28 | "name": "drumnation", 29 | "id": 18486434, 30 | "comment_id": 2822330188, 31 | "created_at": "2025-04-22T19:51:09Z", 32 | "repoId": 840056306, 33 | "pullRequestNo": 389 34 | }, 35 | { 36 | "name": "jackaldenryan", 37 | "id": 61809814, 38 | "comment_id": 2845356793, 39 | "created_at": "2025-05-01T17:51:11Z", 40 | "repoId": 840056306, 41 | "pullRequestNo": 429 42 | }, 43 | { 44 | "name": "t41372", 45 | "id": 36402030, 46 | "comment_id": 2849035400, 47 | "created_at": "2025-05-04T06:24:37Z", 48 | "repoId": 840056306, 49 | "pullRequestNo": 438 50 | }, 51 | { 52 | "name": "markalosey", 53 | "id": 1949914, 54 | "comment_id": 2878173826, 55 | "created_at": "2025-05-13T23:27:16Z", 56 | "repoId": 840056306, 57 | "pullRequestNo": 486 58 | }, 59 | { 60 | "name": "adamkatav", 61 | "id": 13109136, 62 | "comment_id": 2887184706, 63 | "created_at": "2025-05-16T16:29:22Z", 64 | "repoId": 840056306, 65 | "pullRequestNo": 493 66 | }, 67 | { 68 | "name": "realugbun", 69 | "id": 74101927, 70 | "comment_id": 2899731784, 71 | "created_at": "2025-05-22T02:36:44Z", 72 | "repoId": 840056306, 73 | "pullRequestNo": 513 74 | }, 75 | { 76 | "name": "dudizimber", 77 | "id": 16744955, 78 | "comment_id": 2912211548, 79 | "created_at": "2025-05-27T11:45:57Z", 80 | "repoId": 840056306, 81 | "pullRequestNo": 525 82 | }, 83 | { 84 | "name": "galshubeli", 85 | "id": 124919062, 86 | "comment_id": 2912289100, 87 | "created_at": "2025-05-27T12:15:03Z", 88 | "repoId": 840056306, 89 | "pullRequestNo": 525 90 | }, 91 | { 92 | "name": "TheEpTic", 93 | "id": 326774, 94 | "comment_id": 2917970901, 95 | "created_at": "2025-05-29T01:26:54Z", 96 | "repoId": 840056306, 97 | "pullRequestNo": 541 98 | } 99 | ] 100 | } -------------------------------------------------------------------------------- /tests/cross_encoder/test_bge_reranker_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import pytest 18 | 19 | from graphiti_core.cross_encoder.bge_reranker_client import BGERerankerClient 20 | 21 | pytestmark = pytest.mark.integration 22 | 23 | 24 | @pytest.fixture 25 | def client(): 26 | return BGERerankerClient() 27 | 28 | 29 | @pytest.mark.asyncio 30 | @pytest.mark.integration 31 | async def test_rank_basic_functionality(client): 32 | query = 'What is the capital of France?' 33 | passages = [ 34 | 'Paris is the capital and most populous city of France.', 35 | 'London is the capital city of England and the United Kingdom.', 36 | 'Berlin is the capital and largest city of Germany.', 37 | ] 38 | 39 | ranked_passages = await client.rank(query, passages) 40 | 41 | # Check if the output is a list of tuples 42 | assert isinstance(ranked_passages, list) 43 | assert all(isinstance(item, tuple) for item in ranked_passages) 44 | 45 | # Check if the output has the correct length 46 | assert len(ranked_passages) == len(passages) 47 | 48 | # Check if the scores are floats and passages are strings 49 | for passage, score in ranked_passages: 50 | assert isinstance(passage, str) 51 | assert isinstance(score, float) 52 | 53 | # Check if the results are sorted in descending order 54 | scores = [score for _, score in ranked_passages] 55 | assert scores == sorted(scores, reverse=True) 56 | 57 | 58 | @pytest.mark.asyncio 59 | @pytest.mark.integration 60 | async def test_rank_empty_input(client): 61 | query = 'Empty test' 62 | passages = [] 63 | 64 | ranked_passages = await client.rank(query, passages) 65 | 66 | # Check if the output is an empty list 67 | assert ranked_passages == [] 68 | 69 | 70 | @pytest.mark.asyncio 71 | @pytest.mark.integration 72 | async def test_rank_single_passage(client): 73 | query = 'Test query' 74 | passages = ['Single test passage'] 75 | 76 | ranked_passages = await client.rank(query, passages) 77 | 78 | # Check if the output has one item 79 | assert len(ranked_passages) == 1 80 | 81 | # Check if the passage is correct and the score is a float 82 | assert ranked_passages[0][0] == passages[0] 83 | assert isinstance(ranked_passages[0][1], float) 84 | -------------------------------------------------------------------------------- /tests/embedder/embedder_fixtures.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | 18 | def create_embedding_values(multiplier: float = 0.1, dimension: int = 1536) -> list[float]: 19 | """Create embedding values with the specified multiplier and dimension.""" 20 | return [multiplier] * dimension 21 | -------------------------------------------------------------------------------- /tests/embedder/test_gemini.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Generator 18 | from typing import Any 19 | from unittest.mock import AsyncMock, MagicMock, patch 20 | 21 | import pytest 22 | 23 | from graphiti_core.embedder.gemini import ( 24 | DEFAULT_EMBEDDING_MODEL, 25 | GeminiEmbedder, 26 | GeminiEmbedderConfig, 27 | ) 28 | from tests.embedder.embedder_fixtures import create_embedding_values 29 | 30 | 31 | def create_gemini_embedding(multiplier: float = 0.1) -> MagicMock: 32 | """Create a mock Gemini embedding with specified value multiplier.""" 33 | mock_embedding = MagicMock() 34 | mock_embedding.values = create_embedding_values(multiplier) 35 | return mock_embedding 36 | 37 | 38 | @pytest.fixture 39 | def mock_gemini_response() -> MagicMock: 40 | """Create a mock Gemini embeddings response.""" 41 | mock_result = MagicMock() 42 | mock_result.embeddings = [create_gemini_embedding()] 43 | return mock_result 44 | 45 | 46 | @pytest.fixture 47 | def mock_gemini_batch_response() -> MagicMock: 48 | """Create a mock Gemini batch embeddings response.""" 49 | mock_result = MagicMock() 50 | mock_result.embeddings = [ 51 | create_gemini_embedding(0.1), 52 | create_gemini_embedding(0.2), 53 | create_gemini_embedding(0.3), 54 | ] 55 | return mock_result 56 | 57 | 58 | @pytest.fixture 59 | def mock_gemini_client() -> Generator[Any, Any, None]: 60 | """Create a mocked Gemini client.""" 61 | with patch('google.genai.Client') as mock_client: 62 | mock_instance = mock_client.return_value 63 | mock_instance.aio = MagicMock() 64 | mock_instance.aio.models = MagicMock() 65 | mock_instance.aio.models.embed_content = AsyncMock() 66 | yield mock_instance 67 | 68 | 69 | @pytest.fixture 70 | def gemini_embedder(mock_gemini_client: Any) -> GeminiEmbedder: 71 | """Create a GeminiEmbedder with a mocked client.""" 72 | config = GeminiEmbedderConfig(api_key='test_api_key') 73 | client = GeminiEmbedder(config=config) 74 | client.client = mock_gemini_client 75 | return client 76 | 77 | 78 | @pytest.mark.asyncio 79 | async def test_create_calls_api_correctly( 80 | gemini_embedder: GeminiEmbedder, mock_gemini_client: Any, mock_gemini_response: MagicMock 81 | ) -> None: 82 | """Test that create method correctly calls the API and processes the response.""" 83 | # Setup 84 | mock_gemini_client.aio.models.embed_content.return_value = mock_gemini_response 85 | 86 | # Call method 87 | result = await gemini_embedder.create('Test input') 88 | 89 | # Verify API is called with correct parameters 90 | mock_gemini_client.aio.models.embed_content.assert_called_once() 91 | _, kwargs = mock_gemini_client.aio.models.embed_content.call_args 92 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 93 | assert kwargs['contents'] == ['Test input'] 94 | 95 | # Verify result is processed correctly 96 | assert result == mock_gemini_response.embeddings[0].values 97 | 98 | 99 | @pytest.mark.asyncio 100 | async def test_create_batch_processes_multiple_inputs( 101 | gemini_embedder: GeminiEmbedder, mock_gemini_client: Any, mock_gemini_batch_response: MagicMock 102 | ) -> None: 103 | """Test that create_batch method correctly processes multiple inputs.""" 104 | # Setup 105 | mock_gemini_client.aio.models.embed_content.return_value = mock_gemini_batch_response 106 | input_batch = ['Input 1', 'Input 2', 'Input 3'] 107 | 108 | # Call method 109 | result = await gemini_embedder.create_batch(input_batch) 110 | 111 | # Verify API is called with correct parameters 112 | mock_gemini_client.aio.models.embed_content.assert_called_once() 113 | _, kwargs = mock_gemini_client.aio.models.embed_content.call_args 114 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 115 | assert kwargs['contents'] == input_batch 116 | 117 | # Verify all results are processed correctly 118 | assert len(result) == 3 119 | assert result == [ 120 | mock_gemini_batch_response.embeddings[0].values, 121 | mock_gemini_batch_response.embeddings[1].values, 122 | mock_gemini_batch_response.embeddings[2].values, 123 | ] 124 | 125 | 126 | if __name__ == '__main__': 127 | pytest.main(['-xvs', __file__]) 128 | -------------------------------------------------------------------------------- /tests/embedder/test_openai.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Generator 18 | from typing import Any 19 | from unittest.mock import AsyncMock, MagicMock, patch 20 | 21 | import pytest 22 | 23 | from graphiti_core.embedder.openai import ( 24 | DEFAULT_EMBEDDING_MODEL, 25 | OpenAIEmbedder, 26 | OpenAIEmbedderConfig, 27 | ) 28 | from tests.embedder.embedder_fixtures import create_embedding_values 29 | 30 | 31 | def create_openai_embedding(multiplier: float = 0.1) -> MagicMock: 32 | """Create a mock OpenAI embedding with specified value multiplier.""" 33 | mock_embedding = MagicMock() 34 | mock_embedding.embedding = create_embedding_values(multiplier) 35 | return mock_embedding 36 | 37 | 38 | @pytest.fixture 39 | def mock_openai_response() -> MagicMock: 40 | """Create a mock OpenAI embeddings response.""" 41 | mock_result = MagicMock() 42 | mock_result.data = [create_openai_embedding()] 43 | return mock_result 44 | 45 | 46 | @pytest.fixture 47 | def mock_openai_batch_response() -> MagicMock: 48 | """Create a mock OpenAI batch embeddings response.""" 49 | mock_result = MagicMock() 50 | mock_result.data = [ 51 | create_openai_embedding(0.1), 52 | create_openai_embedding(0.2), 53 | create_openai_embedding(0.3), 54 | ] 55 | return mock_result 56 | 57 | 58 | @pytest.fixture 59 | def mock_openai_client() -> Generator[Any, Any, None]: 60 | """Create a mocked OpenAI client.""" 61 | with patch('openai.AsyncOpenAI') as mock_client: 62 | mock_instance = mock_client.return_value 63 | mock_instance.embeddings = MagicMock() 64 | mock_instance.embeddings.create = AsyncMock() 65 | yield mock_instance 66 | 67 | 68 | @pytest.fixture 69 | def openai_embedder(mock_openai_client: Any) -> OpenAIEmbedder: 70 | """Create an OpenAIEmbedder with a mocked client.""" 71 | config = OpenAIEmbedderConfig(api_key='test_api_key') 72 | client = OpenAIEmbedder(config=config) 73 | client.client = mock_openai_client 74 | return client 75 | 76 | 77 | @pytest.mark.asyncio 78 | async def test_create_calls_api_correctly( 79 | openai_embedder: OpenAIEmbedder, mock_openai_client: Any, mock_openai_response: MagicMock 80 | ) -> None: 81 | """Test that create method correctly calls the API and processes the response.""" 82 | # Setup 83 | mock_openai_client.embeddings.create.return_value = mock_openai_response 84 | 85 | # Call method 86 | result = await openai_embedder.create('Test input') 87 | 88 | # Verify API is called with correct parameters 89 | mock_openai_client.embeddings.create.assert_called_once() 90 | _, kwargs = mock_openai_client.embeddings.create.call_args 91 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 92 | assert kwargs['input'] == 'Test input' 93 | 94 | # Verify result is processed correctly 95 | assert result == mock_openai_response.data[0].embedding[: openai_embedder.config.embedding_dim] 96 | 97 | 98 | @pytest.mark.asyncio 99 | async def test_create_batch_processes_multiple_inputs( 100 | openai_embedder: OpenAIEmbedder, mock_openai_client: Any, mock_openai_batch_response: MagicMock 101 | ) -> None: 102 | """Test that create_batch method correctly processes multiple inputs.""" 103 | # Setup 104 | mock_openai_client.embeddings.create.return_value = mock_openai_batch_response 105 | input_batch = ['Input 1', 'Input 2', 'Input 3'] 106 | 107 | # Call method 108 | result = await openai_embedder.create_batch(input_batch) 109 | 110 | # Verify API is called with correct parameters 111 | mock_openai_client.embeddings.create.assert_called_once() 112 | _, kwargs = mock_openai_client.embeddings.create.call_args 113 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 114 | assert kwargs['input'] == input_batch 115 | 116 | # Verify all results are processed correctly 117 | assert len(result) == 3 118 | assert result == [ 119 | mock_openai_batch_response.data[0].embedding[: openai_embedder.config.embedding_dim], 120 | mock_openai_batch_response.data[1].embedding[: openai_embedder.config.embedding_dim], 121 | mock_openai_batch_response.data[2].embedding[: openai_embedder.config.embedding_dim], 122 | ] 123 | 124 | 125 | if __name__ == '__main__': 126 | pytest.main(['-xvs', __file__]) 127 | -------------------------------------------------------------------------------- /tests/embedder/test_voyage.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from collections.abc import Generator 18 | from typing import Any 19 | from unittest.mock import AsyncMock, MagicMock, patch 20 | 21 | import pytest 22 | 23 | from graphiti_core.embedder.voyage import ( 24 | DEFAULT_EMBEDDING_MODEL, 25 | VoyageAIEmbedder, 26 | VoyageAIEmbedderConfig, 27 | ) 28 | from tests.embedder.embedder_fixtures import create_embedding_values 29 | 30 | 31 | @pytest.fixture 32 | def mock_voyageai_response() -> MagicMock: 33 | """Create a mock VoyageAI embeddings response.""" 34 | mock_result = MagicMock() 35 | mock_result.embeddings = [create_embedding_values()] 36 | return mock_result 37 | 38 | 39 | @pytest.fixture 40 | def mock_voyageai_batch_response() -> MagicMock: 41 | """Create a mock VoyageAI batch embeddings response.""" 42 | mock_result = MagicMock() 43 | mock_result.embeddings = [ 44 | create_embedding_values(0.1), 45 | create_embedding_values(0.2), 46 | create_embedding_values(0.3), 47 | ] 48 | return mock_result 49 | 50 | 51 | @pytest.fixture 52 | def mock_voyageai_client() -> Generator[Any, Any, None]: 53 | """Create a mocked VoyageAI client.""" 54 | with patch('voyageai.AsyncClient') as mock_client: 55 | mock_instance = mock_client.return_value 56 | mock_instance.embed = AsyncMock() 57 | yield mock_instance 58 | 59 | 60 | @pytest.fixture 61 | def voyageai_embedder(mock_voyageai_client: Any) -> VoyageAIEmbedder: 62 | """Create a VoyageAIEmbedder with a mocked client.""" 63 | config = VoyageAIEmbedderConfig(api_key='test_api_key') 64 | client = VoyageAIEmbedder(config=config) 65 | client.client = mock_voyageai_client 66 | return client 67 | 68 | 69 | @pytest.mark.asyncio 70 | async def test_create_calls_api_correctly( 71 | voyageai_embedder: VoyageAIEmbedder, 72 | mock_voyageai_client: Any, 73 | mock_voyageai_response: MagicMock, 74 | ) -> None: 75 | """Test that create method correctly calls the API and processes the response.""" 76 | # Setup 77 | mock_voyageai_client.embed.return_value = mock_voyageai_response 78 | 79 | # Call method 80 | result = await voyageai_embedder.create('Test input') 81 | 82 | # Verify API is called with correct parameters 83 | mock_voyageai_client.embed.assert_called_once() 84 | args, kwargs = mock_voyageai_client.embed.call_args 85 | assert args[0] == ['Test input'] 86 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 87 | 88 | # Verify result is processed correctly 89 | expected_result = [ 90 | float(x) 91 | for x in mock_voyageai_response.embeddings[0][: voyageai_embedder.config.embedding_dim] 92 | ] 93 | assert result == expected_result 94 | 95 | 96 | @pytest.mark.asyncio 97 | async def test_create_batch_processes_multiple_inputs( 98 | voyageai_embedder: VoyageAIEmbedder, 99 | mock_voyageai_client: Any, 100 | mock_voyageai_batch_response: MagicMock, 101 | ) -> None: 102 | """Test that create_batch method correctly processes multiple inputs.""" 103 | # Setup 104 | mock_voyageai_client.embed.return_value = mock_voyageai_batch_response 105 | input_batch = ['Input 1', 'Input 2', 'Input 3'] 106 | 107 | # Call method 108 | result = await voyageai_embedder.create_batch(input_batch) 109 | 110 | # Verify API is called with correct parameters 111 | mock_voyageai_client.embed.assert_called_once() 112 | args, kwargs = mock_voyageai_client.embed.call_args 113 | assert args[0] == input_batch 114 | assert kwargs['model'] == DEFAULT_EMBEDDING_MODEL 115 | 116 | # Verify all results are processed correctly 117 | assert len(result) == 3 118 | expected_results = [ 119 | [ 120 | float(x) 121 | for x in mock_voyageai_batch_response.embeddings[0][ 122 | : voyageai_embedder.config.embedding_dim 123 | ] 124 | ], 125 | [ 126 | float(x) 127 | for x in mock_voyageai_batch_response.embeddings[1][ 128 | : voyageai_embedder.config.embedding_dim 129 | ] 130 | ], 131 | [ 132 | float(x) 133 | for x in mock_voyageai_batch_response.embeddings[2][ 134 | : voyageai_embedder.config.embedding_dim 135 | ] 136 | ], 137 | ] 138 | assert result == expected_results 139 | 140 | 141 | if __name__ == '__main__': 142 | pytest.main(['-xvs', __file__]) 143 | -------------------------------------------------------------------------------- /tests/evals/data/longmemeval_data/README.md: -------------------------------------------------------------------------------- 1 | The `longmemeval_oracle` dataset is an open-source dataset that we are using. 2 | We did not create this dataset and it can be found 3 | here: https://huggingface.co/datasets/xiaowu0162/longmemeval/blob/main/longmemeval_oracle. 4 | -------------------------------------------------------------------------------- /tests/evals/eval_cli.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import asyncio 3 | 4 | from tests.evals.eval_e2e_graph_building import build_baseline_graph, eval_graph 5 | 6 | 7 | async def main(): 8 | parser = argparse.ArgumentParser( 9 | description='Run eval_graph and optionally build_baseline_graph from the command line.' 10 | ) 11 | 12 | parser.add_argument( 13 | '--multi-session-count', 14 | type=int, 15 | required=True, 16 | help='Integer representing multi-session count', 17 | ) 18 | parser.add_argument('--session-length', type=int, required=True, help='Length of each session') 19 | parser.add_argument( 20 | '--build-baseline', action='store_true', help='If set, also runs build_baseline_graph' 21 | ) 22 | 23 | args = parser.parse_args() 24 | 25 | # Optionally run the async function 26 | if args.build_baseline: 27 | print('Running build_baseline_graph...') 28 | await build_baseline_graph( 29 | multi_session_count=args.multi_session_count, session_length=args.session_length 30 | ) 31 | 32 | # Always call eval_graph 33 | result = await eval_graph( 34 | multi_session_count=args.multi_session_count, session_length=args.session_length 35 | ) 36 | print('Result of eval_graph:', result) 37 | 38 | 39 | if __name__ == '__main__': 40 | asyncio.run(main()) 41 | -------------------------------------------------------------------------------- /tests/evals/pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | asyncio_default_fixture_loop_scope = function 3 | markers = 4 | integration: marks tests as integration tests -------------------------------------------------------------------------------- /tests/evals/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging 18 | import sys 19 | 20 | 21 | def setup_logging(): 22 | # Create a logger 23 | logger = logging.getLogger() 24 | logger.setLevel(logging.INFO) # Set the logging level to INFO 25 | 26 | # Create console handler and set level to INFO 27 | console_handler = logging.StreamHandler(sys.stdout) 28 | console_handler.setLevel(logging.INFO) 29 | 30 | # Create formatter 31 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 32 | 33 | # Add formatter to console handler 34 | console_handler.setFormatter(formatter) 35 | 36 | # Add console handler to logger 37 | logger.addHandler(console_handler) 38 | 39 | return logger 40 | -------------------------------------------------------------------------------- /tests/helpers_test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import pytest 18 | 19 | from graphiti_core.helpers import lucene_sanitize 20 | 21 | 22 | def test_lucene_sanitize(): 23 | # Call the function with test data 24 | queries = [ 25 | ( 26 | 'This has every escape character + - && || ! ( ) { } [ ] ^ " ~ * ? : \\ /', 27 | '\This has every escape character \+ \- \&\& \|\| \! \( \) \{ \} \[ \] \^ \\" \~ \* \? \: \\\ \/', 28 | ), 29 | ('this has no escape characters', 'this has no escape characters'), 30 | ] 31 | 32 | for query, assert_result in queries: 33 | result = lucene_sanitize(query) 34 | assert assert_result == result 35 | 36 | 37 | if __name__ == '__main__': 38 | pytest.main([__file__]) 39 | -------------------------------------------------------------------------------- /tests/llm_client/test_anthropic_client_int.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | # Running tests: pytest -xvs tests/integrations/test_anthropic_client_int.py 18 | 19 | import os 20 | 21 | import pytest 22 | from pydantic import BaseModel, Field 23 | 24 | from graphiti_core.llm_client.anthropic_client import AnthropicClient 25 | from graphiti_core.prompts.models import Message 26 | 27 | # Skip all tests if no API key is available 28 | pytestmark = pytest.mark.skipif( 29 | 'TEST_ANTHROPIC_API_KEY' not in os.environ, 30 | reason='Anthropic API key not available', 31 | ) 32 | 33 | 34 | # Rename to avoid pytest collection as a test class 35 | class SimpleResponseModel(BaseModel): 36 | """Test response model.""" 37 | 38 | message: str = Field(..., description='A message from the model') 39 | 40 | 41 | @pytest.mark.asyncio 42 | @pytest.mark.integration 43 | async def test_generate_simple_response(): 44 | """Test generating a simple response from the Anthropic API.""" 45 | if 'TEST_ANTHROPIC_API_KEY' not in os.environ: 46 | pytest.skip('Anthropic API key not available') 47 | 48 | client = AnthropicClient() 49 | 50 | messages = [ 51 | Message( 52 | role='user', 53 | content="Respond with a JSON object containing a 'message' field with value 'Hello, world!'", 54 | ) 55 | ] 56 | 57 | try: 58 | response = await client.generate_response(messages, response_model=SimpleResponseModel) 59 | 60 | assert isinstance(response, dict) 61 | assert 'message' in response 62 | assert response['message'] == 'Hello, world!' 63 | except Exception as e: 64 | pytest.skip(f'Test skipped due to Anthropic API error: {str(e)}') 65 | 66 | 67 | @pytest.mark.asyncio 68 | @pytest.mark.integration 69 | async def test_extract_json_from_text(): 70 | """Test the extract_json_from_text method with real data.""" 71 | # We don't need an actual API connection for this test, 72 | # so we can create the client without worrying about the API key 73 | with pytest.MonkeyPatch.context() as monkeypatch: 74 | # Temporarily set an environment variable to avoid API key error 75 | monkeypatch.setenv('ANTHROPIC_API_KEY', 'fake_key_for_testing') 76 | client = AnthropicClient(cache=False) 77 | 78 | # A string with embedded JSON 79 | text = 'Some text before {"message": "Hello, world!"} and after' 80 | 81 | result = client._extract_json_from_text(text) # type: ignore # ignore type check for private method 82 | 83 | assert isinstance(result, dict) 84 | assert 'message' in result 85 | assert result['message'] == 'Hello, world!' 86 | -------------------------------------------------------------------------------- /tests/llm_client/test_client.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | from graphiti_core.llm_client.client import LLMClient 18 | from graphiti_core.llm_client.config import LLMConfig 19 | 20 | 21 | class MockLLMClient(LLMClient): 22 | """Concrete implementation of LLMClient for testing""" 23 | 24 | async def _generate_response(self, messages, response_model=None): 25 | return {'content': 'test'} 26 | 27 | 28 | def test_clean_input(): 29 | client = MockLLMClient(LLMConfig()) 30 | 31 | test_cases = [ 32 | # Basic text should remain unchanged 33 | ('Hello World', 'Hello World'), 34 | # Control characters should be removed 35 | ('Hello\x00World', 'HelloWorld'), 36 | # Newlines, tabs, returns should be preserved 37 | ('Hello\nWorld\tTest\r', 'Hello\nWorld\tTest\r'), 38 | # Invalid Unicode should be removed 39 | ('Hello\udcdeWorld', 'HelloWorld'), 40 | # Zero-width characters should be removed 41 | ('Hello\u200bWorld', 'HelloWorld'), 42 | ('Test\ufeffWord', 'TestWord'), 43 | # Multiple issues combined 44 | ('Hello\x00\u200b\nWorld\udcde', 'Hello\nWorld'), 45 | # Empty string should remain empty 46 | ('', ''), 47 | # Form feed and other control characters from the error case 48 | ('{"edges":[{"relation_typ...\f\x04Hn\\?"}]}', '{"edges":[{"relation_typ...Hn\\?"}]}'), 49 | # More specific control character tests 50 | ('Hello\x0cWorld', 'HelloWorld'), # form feed \f 51 | ('Hello\x04World', 'HelloWorld'), # end of transmission 52 | # Combined JSON-like string with control characters 53 | ('{"test": "value\f\x00\x04"}', '{"test": "value"}'), 54 | ] 55 | 56 | for input_str, expected in test_cases: 57 | assert client._clean_input(input_str) == expected, f'Failed for input: {repr(input_str)}' 58 | -------------------------------------------------------------------------------- /tests/llm_client/test_errors.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | # Running tests: pytest -xvs tests/llm_client/test_errors.py 18 | 19 | import pytest 20 | 21 | from graphiti_core.llm_client.errors import EmptyResponseError, RateLimitError, RefusalError 22 | 23 | 24 | class TestRateLimitError: 25 | """Tests for the RateLimitError class.""" 26 | 27 | def test_default_message(self): 28 | """Test that the default message is set correctly.""" 29 | error = RateLimitError() 30 | assert error.message == 'Rate limit exceeded. Please try again later.' 31 | assert str(error) == 'Rate limit exceeded. Please try again later.' 32 | 33 | def test_custom_message(self): 34 | """Test that a custom message can be set.""" 35 | custom_message = 'Custom rate limit message' 36 | error = RateLimitError(custom_message) 37 | assert error.message == custom_message 38 | assert str(error) == custom_message 39 | 40 | 41 | class TestRefusalError: 42 | """Tests for the RefusalError class.""" 43 | 44 | def test_message_required(self): 45 | """Test that a message is required for RefusalError.""" 46 | with pytest.raises(TypeError): 47 | # Intentionally not providing the required message parameter 48 | RefusalError() # type: ignore 49 | 50 | def test_message_assignment(self): 51 | """Test that the message is assigned correctly.""" 52 | message = 'The LLM refused to respond to this prompt.' 53 | error = RefusalError(message=message) # Add explicit keyword argument 54 | assert error.message == message 55 | assert str(error) == message 56 | 57 | 58 | class TestEmptyResponseError: 59 | """Tests for the EmptyResponseError class.""" 60 | 61 | def test_message_required(self): 62 | """Test that a message is required for EmptyResponseError.""" 63 | with pytest.raises(TypeError): 64 | # Intentionally not providing the required message parameter 65 | EmptyResponseError() # type: ignore 66 | 67 | def test_message_assignment(self): 68 | """Test that the message is assigned correctly.""" 69 | message = 'The LLM returned an empty response.' 70 | error = EmptyResponseError(message=message) # Add explicit keyword argument 71 | assert error.message == message 72 | assert str(error) == message 73 | 74 | 75 | if __name__ == '__main__': 76 | pytest.main(['-v', 'test_errors.py']) 77 | -------------------------------------------------------------------------------- /tests/test_graphiti_int.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import logging 18 | import os 19 | import sys 20 | from datetime import datetime, timezone 21 | 22 | import pytest 23 | from dotenv import load_dotenv 24 | 25 | from graphiti_core.edges import EntityEdge, EpisodicEdge 26 | from graphiti_core.graphiti import Graphiti 27 | from graphiti_core.helpers import semaphore_gather 28 | from graphiti_core.nodes import EntityNode, EpisodicNode 29 | from graphiti_core.search.search_helpers import search_results_to_context_string 30 | 31 | pytestmark = pytest.mark.integration 32 | 33 | pytest_plugins = ('pytest_asyncio',) 34 | 35 | load_dotenv() 36 | 37 | NEO4J_URI = os.getenv('NEO4J_URI') 38 | NEO4j_USER = os.getenv('NEO4J_USER') 39 | NEO4j_PASSWORD = os.getenv('NEO4J_PASSWORD') 40 | 41 | 42 | def setup_logging(): 43 | # Create a logger 44 | logger = logging.getLogger() 45 | logger.setLevel(logging.INFO) # Set the logging level to INFO 46 | 47 | # Create console handler and set level to INFO 48 | console_handler = logging.StreamHandler(sys.stdout) 49 | console_handler.setLevel(logging.INFO) 50 | 51 | # Create formatter 52 | formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') 53 | 54 | # Add formatter to console handler 55 | console_handler.setFormatter(formatter) 56 | 57 | # Add console handler to logger 58 | logger.addHandler(console_handler) 59 | 60 | return logger 61 | 62 | 63 | @pytest.mark.asyncio 64 | async def test_graphiti_init(): 65 | logger = setup_logging() 66 | graphiti = Graphiti(NEO4J_URI, NEO4j_USER, NEO4j_PASSWORD) 67 | 68 | results = await graphiti.search_(query='Who is the user?') 69 | 70 | pretty_results = search_results_to_context_string(results) 71 | 72 | logger.info(pretty_results) 73 | 74 | await graphiti.close() 75 | 76 | 77 | @pytest.mark.asyncio 78 | async def test_graph_integration(): 79 | client = Graphiti(NEO4J_URI, NEO4j_USER, NEO4j_PASSWORD) 80 | embedder = client.embedder 81 | driver = client.driver 82 | 83 | now = datetime.now(timezone.utc) 84 | episode = EpisodicNode( 85 | name='test_episode', 86 | labels=[], 87 | created_at=now, 88 | valid_at=now, 89 | source='message', 90 | source_description='conversation message', 91 | content='Alice likes Bob', 92 | entity_edges=[], 93 | ) 94 | 95 | alice_node = EntityNode( 96 | name='Alice', 97 | labels=[], 98 | created_at=now, 99 | summary='Alice summary', 100 | ) 101 | 102 | bob_node = EntityNode(name='Bob', labels=[], created_at=now, summary='Bob summary') 103 | 104 | episodic_edge_1 = EpisodicEdge( 105 | source_node_uuid=episode.uuid, target_node_uuid=alice_node.uuid, created_at=now 106 | ) 107 | 108 | episodic_edge_2 = EpisodicEdge( 109 | source_node_uuid=episode.uuid, target_node_uuid=bob_node.uuid, created_at=now 110 | ) 111 | 112 | entity_edge = EntityEdge( 113 | source_node_uuid=alice_node.uuid, 114 | target_node_uuid=bob_node.uuid, 115 | created_at=now, 116 | name='likes', 117 | fact='Alice likes Bob', 118 | episodes=[], 119 | expired_at=now, 120 | valid_at=now, 121 | invalid_at=now, 122 | ) 123 | 124 | await entity_edge.generate_embedding(embedder) 125 | 126 | nodes = [episode, alice_node, bob_node] 127 | edges = [episodic_edge_1, episodic_edge_2, entity_edge] 128 | 129 | # test save 130 | await semaphore_gather(*[node.save(driver) for node in nodes]) 131 | await semaphore_gather(*[edge.save(driver) for edge in edges]) 132 | 133 | # test get 134 | assert await EpisodicNode.get_by_uuid(driver, episode.uuid) is not None 135 | assert await EntityNode.get_by_uuid(driver, alice_node.uuid) is not None 136 | assert await EpisodicEdge.get_by_uuid(driver, episodic_edge_1.uuid) is not None 137 | assert await EntityEdge.get_by_uuid(driver, entity_edge.uuid) is not None 138 | 139 | # test delete 140 | await semaphore_gather(*[node.delete(driver) for node in nodes]) 141 | await semaphore_gather(*[edge.delete(driver) for edge in edges]) 142 | -------------------------------------------------------------------------------- /tests/test_node_int.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright 2024, Zep Software, Inc. 3 | 4 | Licensed under the Apache License, Version 2.0 (the "License"); 5 | you may not use this file except in compliance with the License. 6 | You may obtain a copy of the License at 7 | 8 | http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | Unless required by applicable law or agreed to in writing, software 11 | distributed under the License is distributed on an "AS IS" BASIS, 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | See the License for the specific language governing permissions and 14 | limitations under the License. 15 | """ 16 | 17 | import os 18 | from datetime import datetime, timezone 19 | from uuid import uuid4 20 | 21 | import pytest 22 | from neo4j import AsyncGraphDatabase 23 | 24 | from graphiti_core.nodes import ( 25 | CommunityNode, 26 | EntityNode, 27 | EpisodeType, 28 | EpisodicNode, 29 | ) 30 | 31 | NEO4J_URI = os.getenv('NEO4J_URI', 'bolt://localhost:7687') 32 | NEO4J_USER = os.getenv('NEO4J_USER', 'neo4j') 33 | NEO4J_PASSWORD = os.getenv('NEO4J_PASSWORD', 'test') 34 | 35 | 36 | @pytest.fixture 37 | def sample_entity_node(): 38 | return EntityNode( 39 | uuid=str(uuid4()), 40 | name='Test Entity', 41 | group_id='test_group', 42 | labels=['Entity'], 43 | name_embedding=[0.5] * 1024, 44 | summary='Entity Summary', 45 | ) 46 | 47 | 48 | @pytest.fixture 49 | def sample_episodic_node(): 50 | return EpisodicNode( 51 | uuid=str(uuid4()), 52 | name='Episode 1', 53 | group_id='test_group', 54 | source=EpisodeType.text, 55 | source_description='Test source', 56 | content='Some content here', 57 | valid_at=datetime.now(timezone.utc), 58 | ) 59 | 60 | 61 | @pytest.fixture 62 | def sample_community_node(): 63 | return CommunityNode( 64 | uuid=str(uuid4()), 65 | name='Community A', 66 | name_embedding=[0.5] * 1024, 67 | group_id='test_group', 68 | summary='Community summary', 69 | ) 70 | 71 | 72 | @pytest.mark.asyncio 73 | @pytest.mark.integration 74 | async def test_entity_node_save_get_and_delete(sample_entity_node): 75 | neo4j_driver = AsyncGraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) 76 | await sample_entity_node.save(neo4j_driver) 77 | retrieved = await EntityNode.get_by_uuid(neo4j_driver, sample_entity_node.uuid) 78 | assert retrieved.uuid == sample_entity_node.uuid 79 | assert retrieved.name == 'Test Entity' 80 | assert retrieved.group_id == 'test_group' 81 | 82 | await sample_entity_node.delete(neo4j_driver) 83 | 84 | await neo4j_driver.close() 85 | 86 | 87 | @pytest.mark.asyncio 88 | @pytest.mark.integration 89 | async def test_community_node_save_get_and_delete(sample_community_node): 90 | neo4j_driver = AsyncGraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) 91 | 92 | await sample_community_node.save(neo4j_driver) 93 | 94 | retrieved = await CommunityNode.get_by_uuid(neo4j_driver, sample_community_node.uuid) 95 | assert retrieved.uuid == sample_community_node.uuid 96 | assert retrieved.name == 'Community A' 97 | assert retrieved.group_id == 'test_group' 98 | assert retrieved.summary == 'Community summary' 99 | 100 | await sample_community_node.delete(neo4j_driver) 101 | 102 | await neo4j_driver.close() 103 | 104 | 105 | @pytest.mark.asyncio 106 | @pytest.mark.integration 107 | async def test_episodic_node_save_get_and_delete(sample_episodic_node): 108 | neo4j_driver = AsyncGraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) 109 | 110 | await sample_episodic_node.save(neo4j_driver) 111 | 112 | retrieved = await EpisodicNode.get_by_uuid(neo4j_driver, sample_episodic_node.uuid) 113 | assert retrieved.uuid == sample_episodic_node.uuid 114 | assert retrieved.name == 'Episode 1' 115 | assert retrieved.group_id == 'test_group' 116 | assert retrieved.source == EpisodeType.text 117 | assert retrieved.source_description == 'Test source' 118 | assert retrieved.content == 'Some content here' 119 | 120 | await sample_episodic_node.delete(neo4j_driver) 121 | 122 | await neo4j_driver.close() 123 | -------------------------------------------------------------------------------- /tests/utils/maintenance/test_edge_operations.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime, timedelta, timezone 2 | from unittest.mock import MagicMock 3 | 4 | import pytest 5 | 6 | from graphiti_core.edges import EntityEdge 7 | from graphiti_core.nodes import EpisodicNode 8 | 9 | 10 | @pytest.fixture 11 | def mock_llm_client(): 12 | return MagicMock() 13 | 14 | 15 | @pytest.fixture 16 | def mock_extracted_edge(): 17 | return EntityEdge( 18 | source_node_uuid='source_uuid', 19 | target_node_uuid='target_uuid', 20 | name='test_edge', 21 | group_id='group_1', 22 | fact='Test fact', 23 | episodes=['episode_1'], 24 | created_at=datetime.now(timezone.utc), 25 | valid_at=None, 26 | invalid_at=None, 27 | ) 28 | 29 | 30 | @pytest.fixture 31 | def mock_related_edges(): 32 | return [ 33 | EntityEdge( 34 | source_node_uuid='source_uuid_2', 35 | target_node_uuid='target_uuid_2', 36 | name='related_edge', 37 | group_id='group_1', 38 | fact='Related fact', 39 | episodes=['episode_2'], 40 | created_at=datetime.now(timezone.utc) - timedelta(days=1), 41 | valid_at=datetime.now(timezone.utc) - timedelta(days=1), 42 | invalid_at=None, 43 | ) 44 | ] 45 | 46 | 47 | @pytest.fixture 48 | def mock_existing_edges(): 49 | return [ 50 | EntityEdge( 51 | source_node_uuid='source_uuid_3', 52 | target_node_uuid='target_uuid_3', 53 | name='existing_edge', 54 | group_id='group_1', 55 | fact='Existing fact', 56 | episodes=['episode_3'], 57 | created_at=datetime.now(timezone.utc) - timedelta(days=2), 58 | valid_at=datetime.now(timezone.utc) - timedelta(days=2), 59 | invalid_at=None, 60 | ) 61 | ] 62 | 63 | 64 | @pytest.fixture 65 | def mock_current_episode(): 66 | return EpisodicNode( 67 | uuid='episode_1', 68 | content='Current episode content', 69 | valid_at=datetime.now(timezone.utc), 70 | name='Current Episode', 71 | group_id='group_1', 72 | source='message', 73 | source_description='Test source description', 74 | ) 75 | 76 | 77 | @pytest.fixture 78 | def mock_previous_episodes(): 79 | return [ 80 | EpisodicNode( 81 | uuid='episode_2', 82 | content='Previous episode content', 83 | valid_at=datetime.now(timezone.utc) - timedelta(days=1), 84 | name='Previous Episode', 85 | group_id='group_1', 86 | source='message', 87 | source_description='Test source description', 88 | ) 89 | ] 90 | 91 | 92 | # Run the tests 93 | if __name__ == '__main__': 94 | pytest.main([__file__]) 95 | --------------------------------------------------------------------------------