Compare commits
1 commit
main
...
chore_onto
| Author | SHA1 | Date | |
|---|---|---|---|
|
8f993d5304 |
59 changed files with 4287 additions and 5003 deletions
25
.github/workflows/e2e_tests.yml
vendored
25
.github/workflows/e2e_tests.yml
vendored
|
|
@ -237,31 +237,6 @@ jobs:
|
|||
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
|
||||
run: uv run python ./cognee/tests/test_dataset_database_handler.py
|
||||
|
||||
test-dataset-database-deletion:
|
||||
name: Test dataset database deletion in Cognee
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Cognee Setup
|
||||
uses: ./.github/actions/cognee_setup
|
||||
with:
|
||||
python-version: '3.11.x'
|
||||
|
||||
- name: Run dataset databases deletion test
|
||||
env:
|
||||
ENV: 'dev'
|
||||
LLM_MODEL: ${{ secrets.LLM_MODEL }}
|
||||
LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
|
||||
LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
|
||||
LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
|
||||
EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
|
||||
EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
|
||||
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
|
||||
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
|
||||
run: uv run python ./cognee/tests/test_dataset_delete.py
|
||||
|
||||
test-permissions:
|
||||
name: Test permissions with different situations in Cognee
|
||||
runs-on: ubuntu-22.04
|
||||
|
|
|
|||
138
.github/workflows/release.yml
vendored
138
.github/workflows/release.yml
vendored
|
|
@ -1,138 +0,0 @@
|
|||
name: release.yml
|
||||
on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
flavour:
|
||||
required: true
|
||||
default: dev
|
||||
type: choice
|
||||
options:
|
||||
- dev
|
||||
- main
|
||||
description: Dev or Main release
|
||||
|
||||
jobs:
|
||||
release-github:
|
||||
name: Create GitHub Release from ${{ inputs.flavour }}
|
||||
outputs:
|
||||
tag: ${{ steps.create_tag.outputs.tag }}
|
||||
version: ${{ steps.create_tag.outputs.version }}
|
||||
permissions:
|
||||
contents: write
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Check out ${{ inputs.flavour }}
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.flavour }}
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
|
||||
- name: Create and push git tag
|
||||
id: create_tag
|
||||
run: |
|
||||
VERSION="$(uv version --short)"
|
||||
TAG="v${VERSION}"
|
||||
|
||||
echo "Tag to create: ${TAG}"
|
||||
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
echo "tag=${TAG}" >> "$GITHUB_OUTPUT"
|
||||
echo "version=${VERSION}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
git tag "${TAG}"
|
||||
git push origin "${TAG}"
|
||||
|
||||
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
tag_name: ${{ steps.create_tag.outputs.tag }}
|
||||
generate_release_notes: true
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
release-pypi-package:
|
||||
needs: release-github
|
||||
name: Release PyPI Package from ${{ inputs.flavour }}
|
||||
permissions:
|
||||
contents: read
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Check out ${{ inputs.flavour }}
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.flavour }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
|
||||
- name: Install Python
|
||||
run: uv python install
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --locked --all-extras
|
||||
|
||||
- name: Build distributions
|
||||
run: uv build
|
||||
|
||||
- name: Publish ${{ inputs.flavour }} release to PyPI
|
||||
env:
|
||||
UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
|
||||
run: uv publish
|
||||
|
||||
release-docker-image:
|
||||
needs: release-github
|
||||
name: Release Docker Image from ${{ inputs.flavour }}
|
||||
permissions:
|
||||
contents: read
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Check out ${{ inputs.flavour }}
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.flavour }}
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
username: ${{ secrets.DOCKER_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_PASSWORD }}
|
||||
|
||||
- name: Build and push Dev Docker Image
|
||||
if: ${{ inputs.flavour == 'dev' }}
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
platforms: linux/amd64,linux/arm64
|
||||
push: true
|
||||
tags: cognee/cognee:${{ needs.release-github.outputs.version }}
|
||||
labels: |
|
||||
version=${{ needs.release-github.outputs.version }}
|
||||
flavour=${{ inputs.flavour }}
|
||||
cache-from: type=registry,ref=cognee/cognee:buildcache
|
||||
cache-to: type=registry,ref=cognee/cognee:buildcache,mode=max
|
||||
|
||||
- name: Build and push Main Docker Image
|
||||
if: ${{ inputs.flavour == 'main' }}
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
platforms: linux/amd64,linux/arm64
|
||||
push: true
|
||||
tags: |
|
||||
cognee/cognee:${{ needs.release-github.outputs.version }}
|
||||
cognee/cognee:latest
|
||||
labels: |
|
||||
version=${{ needs.release-github.outputs.version }}
|
||||
flavour=${{ inputs.flavour }}
|
||||
cache-from: type=registry,ref=cognee/cognee:buildcache
|
||||
cache-to: type=registry,ref=cognee/cognee:buildcache,mode=max
|
||||
90
.github/workflows/test_llms.yml
vendored
90
.github/workflows/test_llms.yml
vendored
|
|
@ -84,93 +84,3 @@ jobs:
|
|||
EMBEDDING_DIMENSIONS: "3072"
|
||||
EMBEDDING_MAX_TOKENS: "8191"
|
||||
run: uv run python ./examples/python/simple_example.py
|
||||
|
||||
test-bedrock-api-key:
|
||||
name: Run Bedrock API Key Test
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Cognee Setup
|
||||
uses: ./.github/actions/cognee_setup
|
||||
with:
|
||||
python-version: '3.11.x'
|
||||
extra-dependencies: "aws"
|
||||
|
||||
- name: Run Bedrock API Key Simple Example
|
||||
env:
|
||||
LLM_PROVIDER: "bedrock"
|
||||
LLM_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
|
||||
LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
LLM_MAX_TOKENS: "16384"
|
||||
AWS_REGION_NAME: "eu-west-1"
|
||||
EMBEDDING_PROVIDER: "bedrock"
|
||||
EMBEDDING_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
|
||||
EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
|
||||
EMBEDDING_DIMENSIONS: "1024"
|
||||
EMBEDDING_MAX_TOKENS: "8191"
|
||||
run: uv run python ./examples/python/simple_example.py
|
||||
|
||||
test-bedrock-aws-credentials:
|
||||
name: Run Bedrock AWS Credentials Test
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Cognee Setup
|
||||
uses: ./.github/actions/cognee_setup
|
||||
with:
|
||||
python-version: '3.11.x'
|
||||
extra-dependencies: "aws"
|
||||
|
||||
- name: Run Bedrock AWS Credentials Simple Example
|
||||
env:
|
||||
LLM_PROVIDER: "bedrock"
|
||||
LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
LLM_MAX_TOKENS: "16384"
|
||||
AWS_REGION_NAME: "eu-west-1"
|
||||
AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
|
||||
AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
|
||||
EMBEDDING_PROVIDER: "bedrock"
|
||||
EMBEDDING_API_KEY: ${{ secrets.BEDROCK_API_KEY }}
|
||||
EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
|
||||
EMBEDDING_DIMENSIONS: "1024"
|
||||
EMBEDDING_MAX_TOKENS: "8191"
|
||||
run: uv run python ./examples/python/simple_example.py
|
||||
|
||||
test-bedrock-aws-profile:
|
||||
name: Run Bedrock AWS Profile Test
|
||||
runs-on: ubuntu-22.04
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Cognee Setup
|
||||
uses: ./.github/actions/cognee_setup
|
||||
with:
|
||||
python-version: '3.11.x'
|
||||
extra-dependencies: "aws"
|
||||
|
||||
- name: Configure AWS Profile
|
||||
run: |
|
||||
mkdir -p ~/.aws
|
||||
cat > ~/.aws/credentials << EOF
|
||||
[bedrock-test]
|
||||
aws_access_key_id = ${{ secrets.AWS_ACCESS_KEY_ID }}
|
||||
aws_secret_access_key = ${{ secrets.AWS_SECRET_ACCESS_KEY }}
|
||||
EOF
|
||||
|
||||
- name: Run Bedrock AWS Profile Simple Example
|
||||
env:
|
||||
LLM_PROVIDER: "bedrock"
|
||||
LLM_MODEL: "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
LLM_MAX_TOKENS: "16384"
|
||||
AWS_PROFILE_NAME: "bedrock-test"
|
||||
AWS_REGION_NAME: "eu-west-1"
|
||||
EMBEDDING_PROVIDER: "bedrock"
|
||||
EMBEDDING_MODEL: "amazon.titan-embed-text-v2:0"
|
||||
EMBEDDING_DIMENSIONS: "1024"
|
||||
EMBEDDING_MAX_TOKENS: "8191"
|
||||
run: uv run python ./examples/python/simple_example.py
|
||||
|
|
@ -71,7 +71,7 @@ git clone https://github.com/<your-github-username>/cognee.git
|
|||
cd cognee
|
||||
```
|
||||
In case you are working on Vector and Graph Adapters
|
||||
1. Fork the [**cognee-community**](https://github.com/topoteretes/cognee-community) repository
|
||||
1. Fork the [**cognee**](https://github.com/topoteretes/cognee-community) repository
|
||||
2. Clone your fork:
|
||||
```shell
|
||||
git clone https://github.com/<your-github-username>/cognee-community.git
|
||||
|
|
@ -97,21 +97,6 @@ git checkout -b feature/your-feature-name
|
|||
python cognee/cognee/tests/test_library.py
|
||||
```
|
||||
|
||||
### Running Simple Example
|
||||
|
||||
Change .env.example into .env and provide your OPENAI_API_KEY as LLM_API_KEY
|
||||
|
||||
Make sure to run ```shell uv sync ``` in the root cloned folder or set up a virtual environment to run cognee
|
||||
|
||||
```shell
|
||||
python cognee/cognee/examples/python/simple_example.py
|
||||
```
|
||||
or
|
||||
|
||||
```shell
|
||||
uv run python cognee/cognee/examples/python/simple_example.py
|
||||
```
|
||||
|
||||
## 4. 📤 Submitting Changes
|
||||
|
||||
1. Install ruff on your system
|
||||
|
|
|
|||
15
README.md
15
README.md
|
|
@ -66,10 +66,13 @@ Use your data to build personalized and dynamic memory for AI Agents. Cognee let
|
|||
## About Cognee
|
||||
|
||||
Cognee is an open-source tool and platform that transforms your raw data into persistent and dynamic AI memory for Agents. It combines vector search with graph databases to make your documents both searchable by meaning and connected by relationships.
|
||||
Cognee offers default memory creation and search which we describe bellow. But with Cognee you can build your own!
|
||||
|
||||
You can use Cognee in two ways:
|
||||
|
||||
### Cognee Open Source:
|
||||
1. [Self-host Cognee Open Source](https://docs.cognee.ai/getting-started/installation), which stores all data locally by default.
|
||||
2. [Connect to Cognee Cloud](https://platform.cognee.ai/), and get the same OSS stack on managed infrastructure for easier development and productionization.
|
||||
|
||||
### Cognee Open Source (self-hosted):
|
||||
|
||||
- Interconnects any type of data — including past conversations, files, images, and audio transcriptions
|
||||
- Replaces traditional RAG systems with a unified memory layer built on graphs and vectors
|
||||
|
|
@ -77,6 +80,11 @@ Cognee offers default memory creation and search which we describe bellow. But w
|
|||
- Provides Pythonic data pipelines for ingestion from 30+ data sources
|
||||
- Offers high customizability through user-defined tasks, modular pipelines, and built-in search endpoints
|
||||
|
||||
### Cognee Cloud (managed):
|
||||
- Hosted web UI dashboard
|
||||
- Automatic version updates
|
||||
- Resource usage analytics
|
||||
- GDPR compliant, enterprise-grade security
|
||||
|
||||
## Basic Usage & Feature Guide
|
||||
|
||||
|
|
@ -118,7 +126,6 @@ Now, run a minimal pipeline:
|
|||
```python
|
||||
import cognee
|
||||
import asyncio
|
||||
from pprint import pprint
|
||||
|
||||
|
||||
async def main():
|
||||
|
|
@ -136,7 +143,7 @@ async def main():
|
|||
|
||||
# Display the results
|
||||
for result in results:
|
||||
pprint(result)
|
||||
print(result)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
|
|
|||
161
cognee-frontend/package-lock.json
generated
161
cognee-frontend/package-lock.json
generated
|
|
@ -12,7 +12,7 @@
|
|||
"classnames": "^2.5.1",
|
||||
"culori": "^4.0.1",
|
||||
"d3-force-3d": "^3.0.6",
|
||||
"next": "16.1.1",
|
||||
"next": "16.0.4",
|
||||
"react": "^19.2.0",
|
||||
"react-dom": "^19.2.0",
|
||||
"react-force-graph-2d": "^1.27.1",
|
||||
|
|
@ -96,6 +96,7 @@
|
|||
"integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"@babel/code-frame": "^7.27.1",
|
||||
"@babel/generator": "^7.28.5",
|
||||
|
|
@ -1073,9 +1074,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/env": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/env/-/env-16.1.1.tgz",
|
||||
"integrity": "sha512-3oxyM97Sr2PqiVyMyrZUtrtM3jqqFxOQJVuKclDsgj/L728iZt/GyslkN4NwarledZATCenbk4Offjk1hQmaAA==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/env/-/env-16.0.4.tgz",
|
||||
"integrity": "sha512-FDPaVoB1kYhtOz6Le0Jn2QV7RZJ3Ngxzqri7YX4yu3Ini+l5lciR7nA9eNDpKTmDm7LWZtxSju+/CQnwRBn2pA==",
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/@next/eslint-plugin-next": {
|
||||
|
|
@ -1089,9 +1090,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-darwin-arm64": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.1.1.tgz",
|
||||
"integrity": "sha512-JS3m42ifsVSJjSTzh27nW+Igfha3NdBOFScr9C80hHGrWx55pTrVL23RJbqir7k7/15SKlrLHhh/MQzqBBYrQA==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.0.4.tgz",
|
||||
"integrity": "sha512-TN0cfB4HT2YyEio9fLwZY33J+s+vMIgC84gQCOLZOYusW7ptgjIn8RwxQt0BUpoo9XRRVVWEHLld0uhyux1ZcA==",
|
||||
"cpu": [
|
||||
"arm64"
|
||||
],
|
||||
|
|
@ -1105,9 +1106,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-darwin-x64": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.1.1.tgz",
|
||||
"integrity": "sha512-hbyKtrDGUkgkyQi1m1IyD3q4I/3m9ngr+V93z4oKHrPcmxwNL5iMWORvLSGAf2YujL+6HxgVvZuCYZfLfb4bGw==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.0.4.tgz",
|
||||
"integrity": "sha512-XsfI23jvimCaA7e+9f3yMCoVjrny2D11G6H8NCcgv+Ina/TQhKPXB9P4q0WjTuEoyZmcNvPdrZ+XtTh3uPfH7Q==",
|
||||
"cpu": [
|
||||
"x64"
|
||||
],
|
||||
|
|
@ -1121,9 +1122,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-linux-arm64-gnu": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.1.1.tgz",
|
||||
"integrity": "sha512-/fvHet+EYckFvRLQ0jPHJCUI5/B56+2DpI1xDSvi80r/3Ez+Eaa2Yq4tJcRTaB1kqj/HrYKn8Yplm9bNoMJpwQ==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.0.4.tgz",
|
||||
"integrity": "sha512-uo8X7qHDy4YdJUhaoJDMAbL8VT5Ed3lijip2DdBHIB4tfKAvB1XBih6INH2L4qIi4jA0Qq1J0ErxcOocBmUSwg==",
|
||||
"cpu": [
|
||||
"arm64"
|
||||
],
|
||||
|
|
@ -1137,9 +1138,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-linux-arm64-musl": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.1.1.tgz",
|
||||
"integrity": "sha512-MFHrgL4TXNQbBPzkKKur4Fb5ICEJa87HM7fczFs2+HWblM7mMLdco3dvyTI+QmLBU9xgns/EeeINSZD6Ar+oLg==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.0.4.tgz",
|
||||
"integrity": "sha512-pvR/AjNIAxsIz0PCNcZYpH+WmNIKNLcL4XYEfo+ArDi7GsxKWFO5BvVBLXbhti8Coyv3DE983NsitzUsGH5yTw==",
|
||||
"cpu": [
|
||||
"arm64"
|
||||
],
|
||||
|
|
@ -1153,9 +1154,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-linux-x64-gnu": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.1.1.tgz",
|
||||
"integrity": "sha512-20bYDfgOQAPUkkKBnyP9PTuHiJGM7HzNBbuqmD0jiFVZ0aOldz+VnJhbxzjcSabYsnNjMPsE0cyzEudpYxsrUQ==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.0.4.tgz",
|
||||
"integrity": "sha512-2hebpsd5MRRtgqmT7Jj/Wze+wG+ZEXUK2KFFL4IlZ0amEEFADo4ywsifJNeFTQGsamH3/aXkKWymDvgEi+pc2Q==",
|
||||
"cpu": [
|
||||
"x64"
|
||||
],
|
||||
|
|
@ -1169,9 +1170,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-linux-x64-musl": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.1.1.tgz",
|
||||
"integrity": "sha512-9pRbK3M4asAHQRkwaXwu601oPZHghuSC8IXNENgbBSyImHv/zY4K5udBusgdHkvJ/Tcr96jJwQYOll0qU8+fPA==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.0.4.tgz",
|
||||
"integrity": "sha512-pzRXf0LZZ8zMljH78j8SeLncg9ifIOp3ugAFka+Bq8qMzw6hPXOc7wydY7ardIELlczzzreahyTpwsim/WL3Sg==",
|
||||
"cpu": [
|
||||
"x64"
|
||||
],
|
||||
|
|
@ -1185,9 +1186,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-win32-arm64-msvc": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.1.1.tgz",
|
||||
"integrity": "sha512-bdfQkggaLgnmYrFkSQfsHfOhk/mCYmjnrbRCGgkMcoOBZ4n+TRRSLmT/CU5SATzlBJ9TpioUyBW/vWFXTqQRiA==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.0.4.tgz",
|
||||
"integrity": "sha512-7G/yJVzum52B5HOqqbQYX9bJHkN+c4YyZ2AIvEssMHQlbAWOn3iIJjD4sM6ihWsBxuljiTKJovEYlD1K8lCUHw==",
|
||||
"cpu": [
|
||||
"arm64"
|
||||
],
|
||||
|
|
@ -1201,9 +1202,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/@next/swc-win32-x64-msvc": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.1.1.tgz",
|
||||
"integrity": "sha512-Ncwbw2WJ57Al5OX0k4chM68DKhEPlrXBaSXDCi2kPi5f4d8b3ejr3RRJGfKBLrn2YJL5ezNS7w2TZLHSti8CMw==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.0.4.tgz",
|
||||
"integrity": "sha512-0Vy4g8SSeVkuU89g2OFHqGKM4rxsQtihGfenjx2tRckPrge5+gtFnRWGAAwvGXr0ty3twQvcnYjEyOrLHJ4JWA==",
|
||||
"cpu": [
|
||||
"x64"
|
||||
],
|
||||
|
|
@ -1512,66 +1513,6 @@
|
|||
"node": ">=14.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/core": {
|
||||
"version": "1.6.0",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"dependencies": {
|
||||
"@emnapi/wasi-threads": "1.1.0",
|
||||
"tslib": "^2.4.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/runtime": {
|
||||
"version": "1.6.0",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"dependencies": {
|
||||
"tslib": "^2.4.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/wasi-threads": {
|
||||
"version": "1.1.0",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"dependencies": {
|
||||
"tslib": "^2.4.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@napi-rs/wasm-runtime": {
|
||||
"version": "1.0.7",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"dependencies": {
|
||||
"@emnapi/core": "^1.5.0",
|
||||
"@emnapi/runtime": "^1.5.0",
|
||||
"@tybys/wasm-util": "^0.10.1"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@tybys/wasm-util": {
|
||||
"version": "0.10.1",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "MIT",
|
||||
"optional": true,
|
||||
"dependencies": {
|
||||
"tslib": "^2.4.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/tslib": {
|
||||
"version": "2.8.1",
|
||||
"dev": true,
|
||||
"inBundle": true,
|
||||
"license": "0BSD",
|
||||
"optional": true
|
||||
},
|
||||
"node_modules/@tailwindcss/oxide-win32-arm64-msvc": {
|
||||
"version": "4.1.17",
|
||||
"resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.1.17.tgz",
|
||||
|
|
@ -1681,6 +1622,7 @@
|
|||
"integrity": "sha512-MWtvHrGZLFttgeEj28VXHxpmwYbor/ATPYbBfSFZEIRK0ecCFLl2Qo55z52Hss+UV9CRN7trSeq1zbgx7YDWWg==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"csstype": "^3.2.2"
|
||||
}
|
||||
|
|
@ -1748,6 +1690,7 @@
|
|||
"integrity": "sha512-jCzKdm/QK0Kg4V4IK/oMlRZlY+QOcdjv89U2NgKHZk1CYTj82/RVSx1mV/0gqCVMJ/DA+Zf/S4NBWNF8GQ+eqQ==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"@typescript-eslint/scope-manager": "8.48.0",
|
||||
"@typescript-eslint/types": "8.48.0",
|
||||
|
|
@ -2256,6 +2199,7 @@
|
|||
"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"bin": {
|
||||
"acorn": "bin/acorn"
|
||||
},
|
||||
|
|
@ -2547,6 +2491,7 @@
|
|||
"version": "2.8.31",
|
||||
"resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.8.31.tgz",
|
||||
"integrity": "sha512-a28v2eWrrRWPpJSzxc+mKwm0ZtVx/G8SepdQZDArnXYU/XS+IF6mp8aB/4E+hH1tyGCoDo3KlUCdlSxGDsRkAw==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"bin": {
|
||||
"baseline-browser-mapping": "dist/cli.js"
|
||||
|
|
@ -2606,6 +2551,7 @@
|
|||
}
|
||||
],
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"baseline-browser-mapping": "^2.8.25",
|
||||
"caniuse-lite": "^1.0.30001754",
|
||||
|
|
@ -2950,6 +2896,7 @@
|
|||
"resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz",
|
||||
"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
|
||||
"license": "ISC",
|
||||
"peer": true,
|
||||
"engines": {
|
||||
"node": ">=12"
|
||||
}
|
||||
|
|
@ -3425,6 +3372,7 @@
|
|||
"integrity": "sha512-BhHmn2yNOFA9H9JmmIVKJmd288g9hrVRDkdoIgRCRuSySRUHH7r/DI6aAXW9T1WwUuY3DFgrcaqB+deURBLR5g==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"@eslint-community/eslint-utils": "^4.8.0",
|
||||
"@eslint-community/regexpp": "^4.12.1",
|
||||
|
|
@ -5463,14 +5411,14 @@
|
|||
"license": "MIT"
|
||||
},
|
||||
"node_modules/next": {
|
||||
"version": "16.1.1",
|
||||
"resolved": "https://registry.npmjs.org/next/-/next-16.1.1.tgz",
|
||||
"integrity": "sha512-QI+T7xrxt1pF6SQ/JYFz95ro/mg/1Znk5vBebsWwbpejj1T0A23hO7GYEaVac9QUOT2BIMiuzm0L99ooq7k0/w==",
|
||||
"version": "16.0.4",
|
||||
"resolved": "https://registry.npmjs.org/next/-/next-16.0.4.tgz",
|
||||
"integrity": "sha512-vICcxKusY8qW7QFOzTvnRL1ejz2ClTqDKtm1AcUjm2mPv/lVAdgpGNsftsPRIDJOXOjRQO68i1dM8Lp8GZnqoA==",
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"@next/env": "16.1.1",
|
||||
"@next/env": "16.0.4",
|
||||
"@swc/helpers": "0.5.15",
|
||||
"baseline-browser-mapping": "^2.8.3",
|
||||
"caniuse-lite": "^1.0.30001579",
|
||||
"postcss": "8.4.31",
|
||||
"styled-jsx": "5.1.6"
|
||||
|
|
@ -5482,14 +5430,14 @@
|
|||
"node": ">=20.9.0"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
"@next/swc-darwin-arm64": "16.1.1",
|
||||
"@next/swc-darwin-x64": "16.1.1",
|
||||
"@next/swc-linux-arm64-gnu": "16.1.1",
|
||||
"@next/swc-linux-arm64-musl": "16.1.1",
|
||||
"@next/swc-linux-x64-gnu": "16.1.1",
|
||||
"@next/swc-linux-x64-musl": "16.1.1",
|
||||
"@next/swc-win32-arm64-msvc": "16.1.1",
|
||||
"@next/swc-win32-x64-msvc": "16.1.1",
|
||||
"@next/swc-darwin-arm64": "16.0.4",
|
||||
"@next/swc-darwin-x64": "16.0.4",
|
||||
"@next/swc-linux-arm64-gnu": "16.0.4",
|
||||
"@next/swc-linux-arm64-musl": "16.0.4",
|
||||
"@next/swc-linux-x64-gnu": "16.0.4",
|
||||
"@next/swc-linux-x64-musl": "16.0.4",
|
||||
"@next/swc-win32-arm64-msvc": "16.0.4",
|
||||
"@next/swc-win32-x64-msvc": "16.0.4",
|
||||
"sharp": "^0.34.4"
|
||||
},
|
||||
"peerDependencies": {
|
||||
|
|
@ -5861,9 +5809,9 @@
|
|||
}
|
||||
},
|
||||
"node_modules/preact": {
|
||||
"version": "10.28.2",
|
||||
"resolved": "https://registry.npmjs.org/preact/-/preact-10.28.2.tgz",
|
||||
"integrity": "sha512-lbteaWGzGHdlIuiJ0l2Jq454m6kcpI1zNje6d8MlGAFlYvP2GO4ibnat7P74Esfz4sPTdM6UxtTwh/d3pwM9JA==",
|
||||
"version": "10.27.2",
|
||||
"resolved": "https://registry.npmjs.org/preact/-/preact-10.27.2.tgz",
|
||||
"integrity": "sha512-5SYSgFKSyhCbk6SrXyMpqjb5+MQBgfvEKE/OC+PujcY34sOpqtr+0AZQtPYx5IA6VxynQ7rUPCtKzyovpj9Bpg==",
|
||||
"license": "MIT",
|
||||
"funding": {
|
||||
"type": "opencollective",
|
||||
|
|
@ -5927,6 +5875,7 @@
|
|||
"resolved": "https://registry.npmjs.org/react/-/react-19.2.0.tgz",
|
||||
"integrity": "sha512-tmbWg6W31tQLeB5cdIBOicJDJRR2KzXsV7uSK9iNfLWQ5bIZfxuPEHp7M8wiHyHnn0DD1i7w3Zmin0FtkrwoCQ==",
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"engines": {
|
||||
"node": ">=0.10.0"
|
||||
}
|
||||
|
|
@ -5936,6 +5885,7 @@
|
|||
"resolved": "https://registry.npmjs.org/react-dom/-/react-dom-19.2.0.tgz",
|
||||
"integrity": "sha512-UlbRu4cAiGaIewkPyiRGJk0imDN2T3JjieT6spoL2UeSf5od4n5LB/mQ4ejmxhCFT1tYe8IvaFulzynWovsEFQ==",
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"dependencies": {
|
||||
"scheduler": "^0.27.0"
|
||||
},
|
||||
|
|
@ -6674,6 +6624,7 @@
|
|||
"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"engines": {
|
||||
"node": ">=12"
|
||||
},
|
||||
|
|
@ -6836,6 +6787,7 @@
|
|||
"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"peer": true,
|
||||
"bin": {
|
||||
"tsc": "bin/tsc",
|
||||
"tsserver": "bin/tsserver"
|
||||
|
|
@ -7133,6 +7085,7 @@
|
|||
"integrity": "sha512-AvvthqfqrAhNH9dnfmrfKzX5upOdjUVJYFqNSlkmGf64gRaTzlPwz99IHYnVs28qYAybvAlBV+H7pn0saFY4Ig==",
|
||||
"dev": true,
|
||||
"license": "MIT",
|
||||
"peer": true,
|
||||
"funding": {
|
||||
"url": "https://github.com/sponsors/colinhacks"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@
|
|||
"classnames": "^2.5.1",
|
||||
"culori": "^4.0.1",
|
||||
"d3-force-3d": "^3.0.6",
|
||||
"next": "16.1.1",
|
||||
"next": "16.0.4",
|
||||
"react": "^19.2.0",
|
||||
"react-dom": "^19.2.0",
|
||||
"react-force-graph-2d": "^1.27.1",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "cognee-mcp"
|
||||
version = "0.5.0"
|
||||
version = "0.4.0"
|
||||
description = "Cognee MCP server"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
|
@ -9,7 +9,7 @@ dependencies = [
|
|||
# For local cognee repo usage remove comment bellow and add absolute path to cognee. Then run `uv sync --reinstall` in the mcp folder on local cognee changes.
|
||||
#"cognee[postgres,codegraph,gemini,huggingface,docs,neo4j] @ file:/Users/igorilic/Desktop/cognee",
|
||||
# TODO: Remove gemini from optional dependecnies for new Cognee version after 0.3.4
|
||||
"cognee[postgres,docs,neo4j]==0.5.0",
|
||||
"cognee[postgres,docs,neo4j]==0.3.7",
|
||||
"fastmcp>=2.10.0,<3.0.0",
|
||||
"mcp>=1.12.0,<2.0.0",
|
||||
"uv>=0.6.3,<1.0.0",
|
||||
|
|
|
|||
|
|
@ -151,7 +151,7 @@ class CogneeClient:
|
|||
query_type: str,
|
||||
datasets: Optional[List[str]] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
top_k: int = 5,
|
||||
top_k: int = 10,
|
||||
) -> Any:
|
||||
"""
|
||||
Search the knowledge graph.
|
||||
|
|
@ -192,7 +192,7 @@ class CogneeClient:
|
|||
|
||||
with redirect_stdout(sys.stderr):
|
||||
results = await self.cognee.search(
|
||||
query_type=SearchType[query_type.upper()], query_text=query_text, top_k=top_k
|
||||
query_type=SearchType[query_type.upper()], query_text=query_text
|
||||
)
|
||||
return results
|
||||
|
||||
|
|
|
|||
|
|
@ -316,7 +316,7 @@ async def save_interaction(data: str) -> list:
|
|||
|
||||
|
||||
@mcp.tool()
|
||||
async def search(search_query: str, search_type: str, top_k: int = 10) -> list:
|
||||
async def search(search_query: str, search_type: str) -> list:
|
||||
"""
|
||||
Search and query the knowledge graph for insights, information, and connections.
|
||||
|
||||
|
|
@ -389,13 +389,6 @@ async def search(search_query: str, search_type: str, top_k: int = 10) -> list:
|
|||
|
||||
The search_type is case-insensitive and will be converted to uppercase.
|
||||
|
||||
top_k : int, optional
|
||||
Maximum number of results to return (default: 10).
|
||||
Controls the amount of context retrieved from the knowledge graph.
|
||||
- Lower values (3-5): Faster, more focused results
|
||||
- Higher values (10-20): More comprehensive, but slower and more context-heavy
|
||||
Helps manage response size and context window usage in MCP clients.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list
|
||||
|
|
@ -432,32 +425,13 @@ async def search(search_query: str, search_type: str, top_k: int = 10) -> list:
|
|||
|
||||
"""
|
||||
|
||||
async def search_task(search_query: str, search_type: str, top_k: int) -> str:
|
||||
"""
|
||||
Internal task to execute knowledge graph search with result formatting.
|
||||
|
||||
Handles the actual search execution and formats results appropriately
|
||||
for MCP clients based on the search type and execution mode (API vs direct).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
search_query : str
|
||||
The search query in natural language
|
||||
search_type : str
|
||||
Type of search to perform (GRAPH_COMPLETION, CHUNKS, etc.)
|
||||
top_k : int
|
||||
Maximum number of results to return
|
||||
|
||||
Returns
|
||||
-------
|
||||
str
|
||||
Formatted search results as a string, with format depending on search_type
|
||||
"""
|
||||
async def search_task(search_query: str, search_type: str) -> str:
|
||||
"""Search the knowledge graph"""
|
||||
# NOTE: MCP uses stdout to communicate, we must redirect all output
|
||||
# going to stdout ( like the print function ) to stderr.
|
||||
with redirect_stdout(sys.stderr):
|
||||
search_results = await cognee_client.search(
|
||||
query_text=search_query, query_type=search_type, top_k=top_k
|
||||
query_text=search_query, query_type=search_type
|
||||
)
|
||||
|
||||
# Handle different result formats based on API vs direct mode
|
||||
|
|
@ -491,7 +465,7 @@ async def search(search_query: str, search_type: str, top_k: int = 10) -> list:
|
|||
else:
|
||||
return str(search_results)
|
||||
|
||||
search_results = await search_task(search_query, search_type, top_k)
|
||||
search_results = await search_task(search_query, search_type)
|
||||
return [types.TextContent(type="text", text=search_results)]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
Test client for Cognee MCP Server functionality.
|
||||
|
||||
This script tests all the tools and functions available in the Cognee MCP server,
|
||||
including cognify, search, prune, status checks, and utility functions.
|
||||
including cognify, codify, search, prune, status checks, and utility functions.
|
||||
|
||||
Usage:
|
||||
# Set your OpenAI API key first
|
||||
|
|
@ -23,7 +23,6 @@ import tempfile
|
|||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from cognee.shared.logging_utils import setup_logging
|
||||
from logging import ERROR, INFO
|
||||
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
|
@ -36,7 +35,7 @@ from src.server import (
|
|||
load_class,
|
||||
)
|
||||
|
||||
# Set timeout for cognify to complete in
|
||||
# Set timeout for cognify/codify to complete in
|
||||
TIMEOUT = 5 * 60 # 5 min in seconds
|
||||
|
||||
|
||||
|
|
@ -152,9 +151,12 @@ DEBUG = True
|
|||
|
||||
expected_tools = {
|
||||
"cognify",
|
||||
"codify",
|
||||
"search",
|
||||
"prune",
|
||||
"cognify_status",
|
||||
"codify_status",
|
||||
"cognee_add_developer_rules",
|
||||
"list_data",
|
||||
"delete",
|
||||
}
|
||||
|
|
@ -245,6 +247,106 @@ DEBUG = True
|
|||
}
|
||||
print(f"❌ {test_name} test failed: {e}")
|
||||
|
||||
async def test_codify(self):
|
||||
"""Test the codify functionality using MCP client."""
|
||||
print("\n🧪 Testing codify functionality...")
|
||||
try:
|
||||
async with self.mcp_server_session() as session:
|
||||
codify_result = await session.call_tool(
|
||||
"codify", arguments={"repo_path": self.test_repo_dir}
|
||||
)
|
||||
|
||||
start = time.time() # mark the start
|
||||
while True:
|
||||
try:
|
||||
# Wait a moment
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Check if codify processing is finished
|
||||
status_result = await session.call_tool("codify_status", arguments={})
|
||||
if hasattr(status_result, "content") and status_result.content:
|
||||
status_text = (
|
||||
status_result.content[0].text
|
||||
if status_result.content
|
||||
else str(status_result)
|
||||
)
|
||||
else:
|
||||
status_text = str(status_result)
|
||||
|
||||
if str(PipelineRunStatus.DATASET_PROCESSING_COMPLETED) in status_text:
|
||||
break
|
||||
elif time.time() - start > TIMEOUT:
|
||||
raise TimeoutError("Codify did not complete in 5min")
|
||||
except DatabaseNotCreatedError:
|
||||
if time.time() - start > TIMEOUT:
|
||||
raise TimeoutError("Database was not created in 5min")
|
||||
|
||||
self.test_results["codify"] = {
|
||||
"status": "PASS",
|
||||
"result": codify_result,
|
||||
"message": "Codify executed successfully",
|
||||
}
|
||||
print("✅ Codify test passed")
|
||||
|
||||
except Exception as e:
|
||||
self.test_results["codify"] = {
|
||||
"status": "FAIL",
|
||||
"error": str(e),
|
||||
"message": "Codify test failed",
|
||||
}
|
||||
print(f"❌ Codify test failed: {e}")
|
||||
|
||||
async def test_cognee_add_developer_rules(self):
|
||||
"""Test the cognee_add_developer_rules functionality using MCP client."""
|
||||
print("\n🧪 Testing cognee_add_developer_rules functionality...")
|
||||
try:
|
||||
async with self.mcp_server_session() as session:
|
||||
result = await session.call_tool(
|
||||
"cognee_add_developer_rules", arguments={"base_path": self.test_data_dir}
|
||||
)
|
||||
|
||||
start = time.time() # mark the start
|
||||
while True:
|
||||
try:
|
||||
# Wait a moment
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Check if developer rule cognify processing is finished
|
||||
status_result = await session.call_tool("cognify_status", arguments={})
|
||||
if hasattr(status_result, "content") and status_result.content:
|
||||
status_text = (
|
||||
status_result.content[0].text
|
||||
if status_result.content
|
||||
else str(status_result)
|
||||
)
|
||||
else:
|
||||
status_text = str(status_result)
|
||||
|
||||
if str(PipelineRunStatus.DATASET_PROCESSING_COMPLETED) in status_text:
|
||||
break
|
||||
elif time.time() - start > TIMEOUT:
|
||||
raise TimeoutError(
|
||||
"Cognify of developer rules did not complete in 5min"
|
||||
)
|
||||
except DatabaseNotCreatedError:
|
||||
if time.time() - start > TIMEOUT:
|
||||
raise TimeoutError("Database was not created in 5min")
|
||||
|
||||
self.test_results["cognee_add_developer_rules"] = {
|
||||
"status": "PASS",
|
||||
"result": result,
|
||||
"message": "Developer rules addition executed successfully",
|
||||
}
|
||||
print("✅ Developer rules test passed")
|
||||
|
||||
except Exception as e:
|
||||
self.test_results["cognee_add_developer_rules"] = {
|
||||
"status": "FAIL",
|
||||
"error": str(e),
|
||||
"message": "Developer rules test failed",
|
||||
}
|
||||
print(f"❌ Developer rules test failed: {e}")
|
||||
|
||||
async def test_search_functionality(self):
|
||||
"""Test the search functionality with different search types using MCP client."""
|
||||
print("\n🧪 Testing search functionality...")
|
||||
|
|
@ -257,11 +359,7 @@ DEBUG = True
|
|||
# Go through all Cognee search types
|
||||
for search_type in SearchType:
|
||||
# Don't test these search types
|
||||
if search_type in [
|
||||
SearchType.NATURAL_LANGUAGE,
|
||||
SearchType.CYPHER,
|
||||
SearchType.TRIPLET_COMPLETION,
|
||||
]:
|
||||
if search_type in [SearchType.NATURAL_LANGUAGE, SearchType.CYPHER]:
|
||||
break
|
||||
try:
|
||||
async with self.mcp_server_session() as session:
|
||||
|
|
@ -583,6 +681,9 @@ class TestModel:
|
|||
test_name="Cognify2",
|
||||
)
|
||||
|
||||
await self.test_codify()
|
||||
await self.test_cognee_add_developer_rules()
|
||||
|
||||
# Test list_data and delete functionality
|
||||
await self.test_list_data()
|
||||
await self.test_delete()
|
||||
|
|
@ -638,5 +739,7 @@ async def main():
|
|||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from logging import ERROR
|
||||
|
||||
logger = setup_logging(log_level=ERROR)
|
||||
asyncio.run(main())
|
||||
|
|
|
|||
7633
cognee-mcp/uv.lock
generated
7633
cognee-mcp/uv.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -155,7 +155,7 @@ async def add(
|
|||
- LLM_API_KEY: API key for your LLM provider (OpenAI, Anthropic, etc.)
|
||||
|
||||
Optional:
|
||||
- LLM_PROVIDER: "openai" (default), "anthropic", "gemini", "ollama", "mistral", "bedrock"
|
||||
- LLM_PROVIDER: "openai" (default), "anthropic", "gemini", "ollama", "mistral"
|
||||
- LLM_MODEL: Model name (default: "gpt-5-mini")
|
||||
- DEFAULT_USER_EMAIL: Custom default user email
|
||||
- DEFAULT_USER_PASSWORD: Custom default user password
|
||||
|
|
|
|||
|
|
@ -53,7 +53,6 @@ async def cognify(
|
|||
custom_prompt: Optional[str] = None,
|
||||
temporal_cognify: bool = False,
|
||||
data_per_batch: int = 20,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Transform ingested data into a structured knowledge graph.
|
||||
|
|
@ -224,7 +223,6 @@ async def cognify(
|
|||
config=config,
|
||||
custom_prompt=custom_prompt,
|
||||
chunks_per_batch=chunks_per_batch,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
|
||||
|
|
@ -253,7 +251,6 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's
|
|||
config: Config = None,
|
||||
custom_prompt: Optional[str] = None,
|
||||
chunks_per_batch: int = 100,
|
||||
**kwargs,
|
||||
) -> list[Task]:
|
||||
if config is None:
|
||||
ontology_config = get_ontology_env_config()
|
||||
|
|
@ -291,7 +288,6 @@ async def get_default_tasks( # TODO: Find out a better way to do this (Boris's
|
|||
config=config,
|
||||
custom_prompt=custom_prompt,
|
||||
task_config={"batch_size": chunks_per_batch},
|
||||
**kwargs,
|
||||
), # Generate knowledge graphs from the document chunks.
|
||||
Task(
|
||||
summarize_text,
|
||||
|
|
|
|||
|
|
@ -42,9 +42,7 @@ class CognifyPayloadDTO(InDTO):
|
|||
default="", description="Custom prompt for entity extraction and graph generation"
|
||||
)
|
||||
ontology_key: Optional[List[str]] = Field(
|
||||
default=None,
|
||||
examples=[[]],
|
||||
description="Reference to one or more previously uploaded ontologies",
|
||||
default=None, description="Reference to one or more previously uploaded ontologies"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -208,14 +208,14 @@ def get_datasets_router() -> APIRouter:
|
|||
},
|
||||
)
|
||||
|
||||
from cognee.modules.data.methods import delete_dataset
|
||||
from cognee.modules.data.methods import get_dataset, delete_dataset
|
||||
|
||||
dataset = await get_authorized_existing_datasets([dataset_id], "delete", user)
|
||||
dataset = await get_dataset(user.id, dataset_id)
|
||||
|
||||
if dataset is None:
|
||||
raise DatasetNotFoundError(message=f"Dataset ({str(dataset_id)}) not found.")
|
||||
|
||||
await delete_dataset(dataset[0])
|
||||
await delete_dataset(dataset)
|
||||
|
||||
@router.delete(
|
||||
"/{dataset_id}/data/{data_id}",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from fastapi import APIRouter, File, Form, UploadFile, Depends, Request
|
||||
from fastapi import APIRouter, File, Form, UploadFile, Depends, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
from typing import Optional, List
|
||||
|
||||
|
|
@ -15,25 +15,28 @@ def get_ontology_router() -> APIRouter:
|
|||
|
||||
@router.post("", response_model=dict)
|
||||
async def upload_ontology(
|
||||
request: Request,
|
||||
ontology_key: str = Form(...),
|
||||
ontology_file: UploadFile = File(...),
|
||||
description: Optional[str] = Form(None),
|
||||
ontology_key: List[str] = Form(...),
|
||||
ontology_file: List[UploadFile] = File(...),
|
||||
descriptions: Optional[List[str]] = Form(None),
|
||||
user: User = Depends(get_authenticated_user),
|
||||
):
|
||||
"""
|
||||
Upload a single ontology file for later use in cognify operations.
|
||||
Upload ontology files with their respective keys for later use in cognify operations.
|
||||
|
||||
Supports both single and multiple file uploads:
|
||||
- Single file: ontology_key=["key"], ontology_file=[file]
|
||||
- Multiple files: ontology_key=["key1", "key2"], ontology_file=[file1, file2]
|
||||
|
||||
## Request Parameters
|
||||
- **ontology_key** (str): User-defined identifier for the ontology.
|
||||
- **ontology_file** (UploadFile): Single OWL format ontology file
|
||||
- **description** (Optional[str]): Optional description for the ontology.
|
||||
- **ontology_key** (List[str]): Repeated field (e.g. ontology_key=foo&ontology_key=bar) of user-defined identifiers
|
||||
- **ontology_file** (List[UploadFile]): OWL format ontology files
|
||||
- **descriptions** (Optional[List[str]]): Repeated optional descriptions aligned with ontology_key
|
||||
|
||||
## Response
|
||||
Returns metadata about the uploaded ontology including key, filename, size, and upload timestamp.
|
||||
Returns metadata about uploaded ontologies including keys, filenames, sizes, and upload timestamps.
|
||||
|
||||
## Error Codes
|
||||
- **400 Bad Request**: Invalid file format, duplicate key, multiple files uploaded
|
||||
- **400 Bad Request**: Invalid file format, duplicate keys, array length mismatches, file size exceeded
|
||||
- **500 Internal Server Error**: File system or processing errors
|
||||
"""
|
||||
send_telemetry(
|
||||
|
|
@ -46,22 +49,8 @@ def get_ontology_router() -> APIRouter:
|
|||
)
|
||||
|
||||
try:
|
||||
# Enforce: exactly one uploaded file for "ontology_file"
|
||||
form = await request.form()
|
||||
uploaded_files = form.getlist("ontology_file")
|
||||
if len(uploaded_files) != 1:
|
||||
raise ValueError("Only one ontology_file is allowed")
|
||||
|
||||
if ontology_key.strip().startswith(("[", "{")):
|
||||
raise ValueError("ontology_key must be a string")
|
||||
if description is not None and description.strip().startswith(("[", "{")):
|
||||
raise ValueError("description must be a string")
|
||||
|
||||
result = await ontology_service.upload_ontology(
|
||||
ontology_key=ontology_key,
|
||||
file=ontology_file,
|
||||
user=user,
|
||||
description=description,
|
||||
results = await ontology_service.upload_ontologies(
|
||||
ontology_key, ontology_file, user, descriptions
|
||||
)
|
||||
|
||||
return {
|
||||
|
|
@ -73,10 +62,9 @@ def get_ontology_router() -> APIRouter:
|
|||
"uploaded_at": result.uploaded_at,
|
||||
"description": result.description,
|
||||
}
|
||||
for result in results
|
||||
]
|
||||
}
|
||||
except ValueError as e:
|
||||
return JSONResponse(status_code=400, content={"error": str(e)})
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
|
|
|
|||
|
|
@ -33,7 +33,6 @@ async def search(
|
|||
session_id: Optional[str] = None,
|
||||
wide_search_top_k: Optional[int] = 100,
|
||||
triplet_distance_penalty: Optional[float] = 3.5,
|
||||
verbose: bool = False,
|
||||
) -> Union[List[SearchResult], CombinedSearchResult]:
|
||||
"""
|
||||
Search and query the knowledge graph for insights, information, and connections.
|
||||
|
|
@ -124,8 +123,6 @@ async def search(
|
|||
|
||||
session_id: Optional session identifier for caching Q&A interactions. Defaults to 'default_session' if None.
|
||||
|
||||
verbose: If True, returns detailed result information including graph representation (when possible).
|
||||
|
||||
Returns:
|
||||
list: Search results in format determined by query_type:
|
||||
|
||||
|
|
@ -207,7 +204,6 @@ async def search(
|
|||
session_id=session_id,
|
||||
wide_search_top_k=wide_search_top_k,
|
||||
triplet_distance_penalty=triplet_distance_penalty,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
return filtered_search_results
|
||||
|
|
|
|||
|
|
@ -1,4 +1,2 @@
|
|||
from .get_or_create_dataset_database import get_or_create_dataset_database
|
||||
from .resolve_dataset_database_connection_info import resolve_dataset_database_connection_info
|
||||
from .get_graph_dataset_database_handler import get_graph_dataset_database_handler
|
||||
from .get_vector_dataset_database_handler import get_vector_dataset_database_handler
|
||||
|
|
|
|||
|
|
@ -1,10 +0,0 @@
|
|||
from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
|
||||
|
||||
|
||||
def get_graph_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler]
|
||||
return handler
|
||||
|
|
@ -1,10 +0,0 @@
|
|||
from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
|
||||
|
||||
|
||||
def get_vector_dataset_database_handler(dataset_database: DatasetDatabase) -> dict:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler]
|
||||
return handler
|
||||
|
|
@ -1,12 +1,24 @@
|
|||
from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
|
||||
get_graph_dataset_database_handler,
|
||||
)
|
||||
from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
|
||||
get_vector_dataset_database_handler,
|
||||
)
|
||||
from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
|
||||
|
||||
|
||||
async def _get_vector_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler]
|
||||
return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database)
|
||||
|
||||
|
||||
async def _get_graph_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler]
|
||||
return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database)
|
||||
|
||||
|
||||
async def resolve_dataset_database_connection_info(
|
||||
dataset_database: DatasetDatabase,
|
||||
) -> DatasetDatabase:
|
||||
|
|
@ -19,12 +31,6 @@ async def resolve_dataset_database_connection_info(
|
|||
Returns:
|
||||
DatasetDatabase instance with resolved connection info
|
||||
"""
|
||||
vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
|
||||
graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
|
||||
dataset_database = await vector_dataset_database_handler[
|
||||
"handler_instance"
|
||||
].resolve_dataset_connection_info(dataset_database)
|
||||
dataset_database = await graph_dataset_database_handler[
|
||||
"handler_instance"
|
||||
].resolve_dataset_connection_info(dataset_database)
|
||||
dataset_database = await _get_vector_db_connection_info(dataset_database)
|
||||
dataset_database = await _get_graph_db_connection_info(dataset_database)
|
||||
return dataset_database
|
||||
|
|
|
|||
|
|
@ -9,8 +9,6 @@ class S3Config(BaseSettings):
|
|||
aws_access_key_id: Optional[str] = None
|
||||
aws_secret_access_key: Optional[str] = None
|
||||
aws_session_token: Optional[str] = None
|
||||
aws_profile_name: Optional[str] = None
|
||||
aws_bedrock_runtime_endpoint: Optional[str] = None
|
||||
model_config = SettingsConfigDict(env_file=".env", extra="allow")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ class LLMGateway:
|
|||
|
||||
@staticmethod
|
||||
def acreate_structured_output(
|
||||
text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> Coroutine:
|
||||
llm_config = get_llm_config()
|
||||
if llm_config.structured_output_framework.upper() == "BAML":
|
||||
|
|
@ -31,10 +31,7 @@ class LLMGateway:
|
|||
|
||||
llm_client = get_llm_client()
|
||||
return llm_client.acreate_structured_output(
|
||||
text_input=text_input,
|
||||
system_prompt=system_prompt,
|
||||
response_model=response_model,
|
||||
**kwargs,
|
||||
text_input=text_input, system_prompt=system_prompt, response_model=response_model
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from cognee.infrastructure.llm.config import (
|
|||
|
||||
|
||||
async def extract_content_graph(
|
||||
content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None, **kwargs
|
||||
content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None
|
||||
):
|
||||
if custom_prompt:
|
||||
system_prompt = custom_prompt
|
||||
|
|
@ -30,7 +30,7 @@ async def extract_content_graph(
|
|||
system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory)
|
||||
|
||||
content_graph = await LLMGateway.acreate_structured_output(
|
||||
content, system_prompt, response_model, **kwargs
|
||||
content, system_prompt, response_model
|
||||
)
|
||||
|
||||
return content_graph
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ class AnthropicAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from a user query.
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
"""Bedrock LLM adapter module."""
|
||||
|
||||
from .adapter import BedrockAdapter
|
||||
|
||||
__all__ = ["BedrockAdapter"]
|
||||
|
|
@ -1,153 +0,0 @@
|
|||
import litellm
|
||||
import instructor
|
||||
from typing import Type
|
||||
from pydantic import BaseModel
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
from instructor.exceptions import InstructorRetryException
|
||||
|
||||
from cognee.infrastructure.llm.LLMGateway import LLMGateway
|
||||
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
|
||||
LLMInterface,
|
||||
)
|
||||
from cognee.infrastructure.llm.exceptions import (
|
||||
ContentPolicyFilterError,
|
||||
MissingSystemPromptPathError,
|
||||
)
|
||||
from cognee.infrastructure.files.storage.s3_config import get_s3_config
|
||||
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
|
||||
rate_limit_async,
|
||||
rate_limit_sync,
|
||||
sleep_and_retry_async,
|
||||
sleep_and_retry_sync,
|
||||
)
|
||||
from cognee.modules.observability.get_observe import get_observe
|
||||
|
||||
observe = get_observe()
|
||||
|
||||
|
||||
class BedrockAdapter(LLMInterface):
|
||||
"""
|
||||
Adapter for AWS Bedrock API with support for three authentication methods:
|
||||
1. API Key (Bearer Token)
|
||||
2. AWS Credentials (access key + secret key)
|
||||
3. AWS Profile (boto3 credential chain)
|
||||
"""
|
||||
|
||||
name = "Bedrock"
|
||||
model: str
|
||||
api_key: str
|
||||
default_instructor_mode = "json_schema_mode"
|
||||
|
||||
MAX_RETRIES = 5
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
api_key: str = None,
|
||||
max_completion_tokens: int = 16384,
|
||||
streaming: bool = False,
|
||||
instructor_mode: str = None,
|
||||
):
|
||||
self.instructor_mode = instructor_mode if instructor_mode else self.default_instructor_mode
|
||||
|
||||
self.aclient = instructor.from_litellm(
|
||||
litellm.acompletion, mode=instructor.Mode(self.instructor_mode)
|
||||
)
|
||||
self.client = instructor.from_litellm(litellm.completion)
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.max_completion_tokens = max_completion_tokens
|
||||
self.streaming = streaming
|
||||
|
||||
def _create_bedrock_request(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> dict:
|
||||
"""Create Bedrock request with authentication."""
|
||||
|
||||
request_params = {
|
||||
"model": self.model,
|
||||
"custom_llm_provider": "bedrock",
|
||||
"drop_params": True,
|
||||
"messages": [
|
||||
{"role": "user", "content": text_input},
|
||||
{"role": "system", "content": system_prompt},
|
||||
],
|
||||
"response_model": response_model,
|
||||
"max_retries": self.MAX_RETRIES,
|
||||
"max_completion_tokens": self.max_completion_tokens,
|
||||
"stream": self.streaming,
|
||||
}
|
||||
|
||||
s3_config = get_s3_config()
|
||||
|
||||
# Add authentication parameters
|
||||
if self.api_key:
|
||||
request_params["api_key"] = self.api_key
|
||||
elif s3_config.aws_access_key_id and s3_config.aws_secret_access_key:
|
||||
request_params["aws_access_key_id"] = s3_config.aws_access_key_id
|
||||
request_params["aws_secret_access_key"] = s3_config.aws_secret_access_key
|
||||
if s3_config.aws_session_token:
|
||||
request_params["aws_session_token"] = s3_config.aws_session_token
|
||||
elif s3_config.aws_profile_name:
|
||||
request_params["aws_profile_name"] = s3_config.aws_profile_name
|
||||
|
||||
if s3_config.aws_region:
|
||||
request_params["aws_region_name"] = s3_config.aws_region
|
||||
|
||||
# Add optional parameters
|
||||
if s3_config.aws_bedrock_runtime_endpoint:
|
||||
request_params["aws_bedrock_runtime_endpoint"] = s3_config.aws_bedrock_runtime_endpoint
|
||||
|
||||
return request_params
|
||||
|
||||
@observe(as_type="generation")
|
||||
@sleep_and_retry_async()
|
||||
@rate_limit_async
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""Generate structured output from AWS Bedrock API."""
|
||||
|
||||
try:
|
||||
request_params = self._create_bedrock_request(text_input, system_prompt, response_model)
|
||||
return await self.aclient.chat.completions.create(**request_params)
|
||||
|
||||
except (
|
||||
ContentPolicyViolationError,
|
||||
InstructorRetryException,
|
||||
) as error:
|
||||
if (
|
||||
isinstance(error, InstructorRetryException)
|
||||
and "content management policy" not in str(error).lower()
|
||||
):
|
||||
raise error
|
||||
|
||||
raise ContentPolicyFilterError(
|
||||
f"The provided input contains content that is not aligned with our content policy: {text_input}"
|
||||
)
|
||||
|
||||
@observe
|
||||
@sleep_and_retry_sync()
|
||||
@rate_limit_sync
|
||||
def create_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""Generate structured output from AWS Bedrock API (synchronous)."""
|
||||
|
||||
request_params = self._create_bedrock_request(text_input, system_prompt, response_model)
|
||||
return self.client.chat.completions.create(**request_params)
|
||||
|
||||
def show_prompt(self, text_input: str, system_prompt: str) -> str:
|
||||
"""Format and display the prompt for a user query."""
|
||||
if not text_input:
|
||||
text_input = "No user input provided."
|
||||
if not system_prompt:
|
||||
raise MissingSystemPromptPathError()
|
||||
system_prompt = LLMGateway.read_query_prompt(system_prompt)
|
||||
|
||||
formatted_prompt = (
|
||||
f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
|
||||
if system_prompt
|
||||
else None
|
||||
)
|
||||
return formatted_prompt
|
||||
|
|
@ -80,7 +80,7 @@ class GeminiAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from a user query.
|
||||
|
|
|
|||
|
|
@ -80,7 +80,7 @@ class GenericAPIAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from a user query.
|
||||
|
|
|
|||
|
|
@ -24,7 +24,6 @@ class LLMProvider(Enum):
|
|||
- CUSTOM: Represents a custom provider option.
|
||||
- GEMINI: Represents the Gemini provider.
|
||||
- MISTRAL: Represents the Mistral AI provider.
|
||||
- BEDROCK: Represents the AWS Bedrock provider.
|
||||
"""
|
||||
|
||||
OPENAI = "openai"
|
||||
|
|
@ -33,7 +32,6 @@ class LLMProvider(Enum):
|
|||
CUSTOM = "custom"
|
||||
GEMINI = "gemini"
|
||||
MISTRAL = "mistral"
|
||||
BEDROCK = "bedrock"
|
||||
|
||||
|
||||
def get_llm_client(raise_api_key_error: bool = True):
|
||||
|
|
@ -156,7 +154,7 @@ def get_llm_client(raise_api_key_error: bool = True):
|
|||
)
|
||||
|
||||
elif provider == LLMProvider.MISTRAL:
|
||||
if llm_config.llm_api_key is None and raise_api_key_error:
|
||||
if llm_config.llm_api_key is None:
|
||||
raise LLMAPIKeyNotSetError()
|
||||
|
||||
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.mistral.adapter import (
|
||||
|
|
@ -171,21 +169,5 @@ def get_llm_client(raise_api_key_error: bool = True):
|
|||
instructor_mode=llm_config.llm_instructor_mode.lower(),
|
||||
)
|
||||
|
||||
elif provider == LLMProvider.BEDROCK:
|
||||
# if llm_config.llm_api_key is None and raise_api_key_error:
|
||||
# raise LLMAPIKeyNotSetError()
|
||||
|
||||
from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.bedrock.adapter import (
|
||||
BedrockAdapter,
|
||||
)
|
||||
|
||||
return BedrockAdapter(
|
||||
model=llm_config.llm_model,
|
||||
api_key=llm_config.llm_api_key,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
streaming=llm_config.llm_streaming,
|
||||
instructor_mode=llm_config.llm_instructor_mode.lower(),
|
||||
)
|
||||
|
||||
else:
|
||||
raise UnsupportedLLMProviderError(provider)
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ class MistralAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from the user query.
|
||||
|
|
|
|||
|
|
@ -76,7 +76,7 @@ class OllamaAPIAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a structured output from the LLM using the provided text and system prompt.
|
||||
|
|
@ -123,7 +123,7 @@ class OllamaAPIAdapter(LLMInterface):
|
|||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||
reraise=True,
|
||||
)
|
||||
async def create_transcript(self, input_file: str, **kwargs) -> str:
|
||||
async def create_transcript(self, input_file: str) -> str:
|
||||
"""
|
||||
Generate an audio transcript from a user query.
|
||||
|
||||
|
|
@ -162,7 +162,7 @@ class OllamaAPIAdapter(LLMInterface):
|
|||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||
reraise=True,
|
||||
)
|
||||
async def transcribe_image(self, input_file: str, **kwargs) -> str:
|
||||
async def transcribe_image(self, input_file: str) -> str:
|
||||
"""
|
||||
Transcribe content from an image using base64 encoding.
|
||||
|
||||
|
|
|
|||
|
|
@ -112,7 +112,7 @@ class OpenAIAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
async def acreate_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from a user query.
|
||||
|
|
@ -154,7 +154,6 @@ class OpenAIAdapter(LLMInterface):
|
|||
api_version=self.api_version,
|
||||
response_model=response_model,
|
||||
max_retries=self.MAX_RETRIES,
|
||||
**kwargs,
|
||||
)
|
||||
except (
|
||||
ContentFilterFinishReasonError,
|
||||
|
|
@ -181,7 +180,6 @@ class OpenAIAdapter(LLMInterface):
|
|||
# api_base=self.fallback_endpoint,
|
||||
response_model=response_model,
|
||||
max_retries=self.MAX_RETRIES,
|
||||
**kwargs,
|
||||
)
|
||||
except (
|
||||
ContentFilterFinishReasonError,
|
||||
|
|
@ -207,7 +205,7 @@ class OpenAIAdapter(LLMInterface):
|
|||
reraise=True,
|
||||
)
|
||||
def create_structured_output(
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel], **kwargs
|
||||
self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
|
||||
) -> BaseModel:
|
||||
"""
|
||||
Generate a response from a user query.
|
||||
|
|
@ -247,7 +245,6 @@ class OpenAIAdapter(LLMInterface):
|
|||
api_version=self.api_version,
|
||||
response_model=response_model,
|
||||
max_retries=self.MAX_RETRIES,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@retry(
|
||||
|
|
@ -257,7 +254,7 @@ class OpenAIAdapter(LLMInterface):
|
|||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||
reraise=True,
|
||||
)
|
||||
async def create_transcript(self, input, **kwargs):
|
||||
async def create_transcript(self, input):
|
||||
"""
|
||||
Generate an audio transcript from a user query.
|
||||
|
||||
|
|
@ -284,7 +281,6 @@ class OpenAIAdapter(LLMInterface):
|
|||
api_base=self.endpoint,
|
||||
api_version=self.api_version,
|
||||
max_retries=self.MAX_RETRIES,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return transcription
|
||||
|
|
@ -296,7 +292,7 @@ class OpenAIAdapter(LLMInterface):
|
|||
before_sleep=before_sleep_log(logger, logging.DEBUG),
|
||||
reraise=True,
|
||||
)
|
||||
async def transcribe_image(self, input, **kwargs) -> BaseModel:
|
||||
async def transcribe_image(self, input) -> BaseModel:
|
||||
"""
|
||||
Generate a transcription of an image from a user query.
|
||||
|
||||
|
|
@ -341,5 +337,4 @@ class OpenAIAdapter(LLMInterface):
|
|||
api_version=self.api_version,
|
||||
max_completion_tokens=300,
|
||||
max_retries=self.MAX_RETRIES,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,10 +5,6 @@ from cognee.context_global_variables import backend_access_control_enabled
|
|||
from cognee.infrastructure.databases.vector import get_vector_engine
|
||||
from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
|
||||
from cognee.infrastructure.databases.relational import get_relational_engine
|
||||
from cognee.infrastructure.databases.utils import (
|
||||
get_graph_dataset_database_handler,
|
||||
get_vector_dataset_database_handler,
|
||||
)
|
||||
from cognee.shared.cache import delete_cache
|
||||
from cognee.modules.users.models import DatasetDatabase
|
||||
from cognee.shared.logging_utils import get_logger
|
||||
|
|
@ -17,13 +13,22 @@ logger = get_logger()
|
|||
|
||||
|
||||
async def prune_graph_databases():
|
||||
async def _prune_graph_db(dataset_database: DatasetDatabase) -> dict:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[
|
||||
dataset_database.graph_dataset_database_handler
|
||||
]
|
||||
return await handler["handler_instance"].delete_dataset(dataset_database)
|
||||
|
||||
db_engine = get_relational_engine()
|
||||
try:
|
||||
dataset_databases = await db_engine.get_all_data_from_table("dataset_database")
|
||||
data = await db_engine.get_all_data_from_table("dataset_database")
|
||||
# Go through each dataset database and delete the graph database
|
||||
for dataset_database in dataset_databases:
|
||||
handler = get_graph_dataset_database_handler(dataset_database)
|
||||
await handler["handler_instance"].delete_dataset(dataset_database)
|
||||
for data_item in data:
|
||||
await _prune_graph_db(data_item)
|
||||
except (OperationalError, EntityNotFoundError) as e:
|
||||
logger.debug(
|
||||
"Skipping pruning of graph DB. Error when accessing dataset_database table: %s",
|
||||
|
|
@ -33,13 +38,22 @@ async def prune_graph_databases():
|
|||
|
||||
|
||||
async def prune_vector_databases():
|
||||
async def _prune_vector_db(dataset_database: DatasetDatabase) -> dict:
|
||||
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
||||
supported_dataset_database_handlers,
|
||||
)
|
||||
|
||||
handler = supported_dataset_database_handlers[
|
||||
dataset_database.vector_dataset_database_handler
|
||||
]
|
||||
return await handler["handler_instance"].delete_dataset(dataset_database)
|
||||
|
||||
db_engine = get_relational_engine()
|
||||
try:
|
||||
dataset_databases = await db_engine.get_all_data_from_table("dataset_database")
|
||||
data = await db_engine.get_all_data_from_table("dataset_database")
|
||||
# Go through each dataset database and delete the vector database
|
||||
for dataset_database in dataset_databases:
|
||||
handler = get_vector_dataset_database_handler(dataset_database)
|
||||
await handler["handler_instance"].delete_dataset(dataset_database)
|
||||
for data_item in data:
|
||||
await _prune_vector_db(data_item)
|
||||
except (OperationalError, EntityNotFoundError) as e:
|
||||
logger.debug(
|
||||
"Skipping pruning of vector DB. Error when accessing dataset_database table: %s",
|
||||
|
|
|
|||
|
|
@ -1,34 +1,8 @@
|
|||
from cognee.modules.users.models import DatasetDatabase
|
||||
from sqlalchemy import select
|
||||
|
||||
from cognee.modules.data.models import Dataset
|
||||
from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
|
||||
get_vector_dataset_database_handler,
|
||||
)
|
||||
from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
|
||||
get_graph_dataset_database_handler,
|
||||
)
|
||||
from cognee.infrastructure.databases.relational import get_relational_engine
|
||||
|
||||
|
||||
async def delete_dataset(dataset: Dataset):
|
||||
db_engine = get_relational_engine()
|
||||
|
||||
async with db_engine.get_async_session() as session:
|
||||
stmt = select(DatasetDatabase).where(
|
||||
DatasetDatabase.dataset_id == dataset.id,
|
||||
)
|
||||
dataset_database: DatasetDatabase = await session.scalar(stmt)
|
||||
if dataset_database:
|
||||
graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
|
||||
vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
|
||||
await graph_dataset_database_handler["handler_instance"].delete_dataset(
|
||||
dataset_database
|
||||
)
|
||||
await vector_dataset_database_handler["handler_instance"].delete_dataset(
|
||||
dataset_database
|
||||
)
|
||||
# TODO: Remove dataset from pipeline_run_status in Data objects related to dataset as well
|
||||
# This blocks recreation of the dataset with the same name and data after deletion as
|
||||
# it's marked as completed and will be just skipped even though it's empty.
|
||||
return await db_engine.delete_entity_by_id(dataset.__tablename__, dataset.id)
|
||||
|
|
|
|||
|
|
@ -15,9 +15,3 @@ async def setup():
|
|||
"""
|
||||
await create_relational_db_and_tables()
|
||||
await create_pgvector_db_and_tables()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
asyncio.run(setup())
|
||||
|
|
|
|||
|
|
@ -49,7 +49,6 @@ async def search(
|
|||
session_id: Optional[str] = None,
|
||||
wide_search_top_k: Optional[int] = 100,
|
||||
triplet_distance_penalty: Optional[float] = 3.5,
|
||||
verbose: bool = False,
|
||||
) -> Union[CombinedSearchResult, List[SearchResult]]:
|
||||
"""
|
||||
|
||||
|
|
@ -141,7 +140,6 @@ async def search(
|
|||
)
|
||||
|
||||
if use_combined_context:
|
||||
# Note: combined context search must always be verbose and return a CombinedSearchResult with graphs info
|
||||
prepared_search_results = await prepare_search_result(
|
||||
search_results[0] if isinstance(search_results, list) else search_results
|
||||
)
|
||||
|
|
@ -175,30 +173,25 @@ async def search(
|
|||
datasets = prepared_search_results["datasets"]
|
||||
|
||||
if only_context:
|
||||
search_result_dict = {
|
||||
"search_result": [context] if context else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
}
|
||||
if verbose:
|
||||
# Include graphs only in verbose mode
|
||||
search_result_dict["graphs"] = graphs
|
||||
|
||||
return_value.append(search_result_dict)
|
||||
return_value.append(
|
||||
{
|
||||
"search_result": [context] if context else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
"graphs": graphs,
|
||||
}
|
||||
)
|
||||
else:
|
||||
search_result_dict = {
|
||||
"search_result": [result] if result else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
}
|
||||
if verbose:
|
||||
# Include graphs only in verbose mode
|
||||
search_result_dict["graphs"] = graphs
|
||||
|
||||
return_value.append(search_result_dict)
|
||||
|
||||
return_value.append(
|
||||
{
|
||||
"search_result": [result] if result else None,
|
||||
"dataset_id": datasets[0].id,
|
||||
"dataset_name": datasets[0].name,
|
||||
"dataset_tenant_id": datasets[0].tenant_id,
|
||||
"graphs": graphs,
|
||||
}
|
||||
)
|
||||
return return_value
|
||||
else:
|
||||
return_value = []
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ class ModelName(Enum):
|
|||
anthropic = "anthropic"
|
||||
gemini = "gemini"
|
||||
mistral = "mistral"
|
||||
bedrock = "bedrock"
|
||||
|
||||
|
||||
class LLMConfig(BaseModel):
|
||||
|
|
@ -78,10 +77,6 @@ def get_settings() -> SettingsDict:
|
|||
"value": "mistral",
|
||||
"label": "Mistral",
|
||||
},
|
||||
{
|
||||
"value": "bedrock",
|
||||
"label": "Bedrock",
|
||||
},
|
||||
]
|
||||
|
||||
return SettingsDict.model_validate(
|
||||
|
|
@ -162,20 +157,6 @@ def get_settings() -> SettingsDict:
|
|||
"label": "Mistral Large 2.1",
|
||||
},
|
||||
],
|
||||
"bedrock": [
|
||||
{
|
||||
"value": "eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"label": "Claude 4.5 Sonnet",
|
||||
},
|
||||
{
|
||||
"value": "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"label": "Claude 4.5 Haiku",
|
||||
},
|
||||
{
|
||||
"value": "eu.amazon.nova-lite-v1:0",
|
||||
"label": "Amazon Nova Lite",
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
vector_db={
|
||||
|
|
|
|||
|
|
@ -92,7 +92,7 @@ async def cognee_network_visualization(graph_data, destination_file_path: str =
|
|||
}
|
||||
links_list.append(link_data)
|
||||
|
||||
html_template = r"""
|
||||
html_template = """
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
|
|
|
|||
|
|
@ -97,7 +97,6 @@ async def extract_graph_from_data(
|
|||
graph_model: Type[BaseModel],
|
||||
config: Config = None,
|
||||
custom_prompt: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> List[DocumentChunk]:
|
||||
"""
|
||||
Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model.
|
||||
|
|
@ -112,7 +111,7 @@ async def extract_graph_from_data(
|
|||
|
||||
chunk_graphs = await asyncio.gather(
|
||||
*[
|
||||
extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt, **kwargs)
|
||||
extract_content_graph(chunk.text, graph_model, custom_prompt=custom_prompt)
|
||||
for chunk in data_chunks
|
||||
]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from typing import AsyncGenerator, Dict, Any, List, Optional
|
||||
from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
|
||||
from cognee.modules.engine.utils import generate_node_id
|
||||
from cognee.shared.logging_utils import get_logger
|
||||
from cognee.modules.graph.utils.convert_node_to_data_point import get_all_subclasses
|
||||
from cognee.infrastructure.engine import DataPoint
|
||||
|
|
@ -156,12 +155,7 @@ def _process_single_triplet(
|
|||
|
||||
embeddable_text = f"{start_node_text}-›{relationship_text}-›{end_node_text}".strip()
|
||||
|
||||
relationship_name = relationship.get("relationship_name", "")
|
||||
triplet_id = generate_node_id(str(start_node_id) + str(relationship_name) + str(end_node_id))
|
||||
|
||||
triplet_obj = Triplet(
|
||||
id=triplet_id, from_node_id=start_node_id, to_node_id=end_node_id, text=embeddable_text
|
||||
)
|
||||
triplet_obj = Triplet(from_node_id=start_node_id, to_node_id=end_node_id, text=embeddable_text)
|
||||
|
||||
return triplet_obj, None
|
||||
|
||||
|
|
|
|||
|
|
@ -148,8 +148,8 @@ class TestCogneeServerStart(unittest.TestCase):
|
|||
headers=headers,
|
||||
files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
|
||||
data={
|
||||
"ontology_key": ontology_key,
|
||||
"description": "Test ontology",
|
||||
"ontology_key": json.dumps([ontology_key]),
|
||||
"description": json.dumps(["Test ontology"]),
|
||||
},
|
||||
)
|
||||
self.assertEqual(ontology_response.status_code, 200)
|
||||
|
|
|
|||
|
|
@ -1,76 +0,0 @@
|
|||
import os
|
||||
import asyncio
|
||||
import pathlib
|
||||
from uuid import UUID
|
||||
|
||||
import cognee
|
||||
from cognee.shared.logging_utils import setup_logging, ERROR
|
||||
from cognee.modules.data.methods.delete_dataset import delete_dataset
|
||||
from cognee.modules.data.methods.get_dataset import get_dataset
|
||||
from cognee.modules.users.methods import get_default_user
|
||||
|
||||
|
||||
async def main():
|
||||
# Set data and system directory paths
|
||||
data_directory_path = str(
|
||||
pathlib.Path(
|
||||
os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_dataset_delete")
|
||||
).resolve()
|
||||
)
|
||||
cognee.config.data_root_directory(data_directory_path)
|
||||
cognee_directory_path = str(
|
||||
pathlib.Path(
|
||||
os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_dataset_delete")
|
||||
).resolve()
|
||||
)
|
||||
cognee.config.system_root_directory(cognee_directory_path)
|
||||
|
||||
# Create a clean slate for cognee -- reset data and system state
|
||||
print("Resetting cognee data...")
|
||||
await cognee.prune.prune_data()
|
||||
await cognee.prune.prune_system(metadata=True)
|
||||
print("Data reset complete.\n")
|
||||
|
||||
# cognee knowledge graph will be created based on this text
|
||||
text = """
|
||||
Natural language processing (NLP) is an interdisciplinary
|
||||
subfield of computer science and information retrieval.
|
||||
"""
|
||||
|
||||
# Add the text, and make it available for cognify
|
||||
await cognee.add(text, "nlp_dataset")
|
||||
await cognee.add("Quantum computing is the study of quantum computers.", "quantum_dataset")
|
||||
|
||||
# Use LLMs and cognee to create knowledge graph
|
||||
ret_val = await cognee.cognify()
|
||||
user = await get_default_user()
|
||||
|
||||
for val in ret_val:
|
||||
dataset_id = str(val)
|
||||
vector_db_path = os.path.join(
|
||||
cognee_directory_path, "databases", str(user.id), dataset_id + ".lance.db"
|
||||
)
|
||||
graph_db_path = os.path.join(
|
||||
cognee_directory_path, "databases", str(user.id), dataset_id + ".pkl"
|
||||
)
|
||||
|
||||
# Check if databases are properly created and exist before deletion
|
||||
assert os.path.exists(graph_db_path), "Graph database file not found."
|
||||
assert os.path.exists(vector_db_path), "Vector database file not found."
|
||||
|
||||
dataset = await get_dataset(user_id=user.id, dataset_id=UUID(dataset_id))
|
||||
await delete_dataset(dataset)
|
||||
|
||||
# Confirm databases have been deleted
|
||||
assert not os.path.exists(graph_db_path), "Graph database file found."
|
||||
assert not os.path.exists(vector_db_path), "Vector database file found."
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger = setup_logging(log_level=ERROR)
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
loop.run_until_complete(main())
|
||||
finally:
|
||||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||||
138
cognee/tests/unit/api/Ontologies_test.py
Normal file
138
cognee/tests/unit/api/Ontologies_test.py
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
import io
|
||||
import json
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi import UploadFile
|
||||
|
||||
from cognee.api.v1.ontologies.ontologies import OntologyService
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_single_ontology_creates_metadata(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
service = OntologyService()
|
||||
user = SimpleNamespace(id="ontology-user")
|
||||
file_content = b"<rdf:RDF>Ontology content</rdf:RDF>"
|
||||
ontology_file = UploadFile(filename="animals.owl", file=io.BytesIO(file_content))
|
||||
|
||||
result = await service.upload_ontology(
|
||||
ontology_key="animals",
|
||||
file=ontology_file,
|
||||
user=user,
|
||||
description="Animal relationships",
|
||||
)
|
||||
|
||||
assert result.ontology_key == "animals"
|
||||
assert result.filename == "animals.owl"
|
||||
assert result.size_bytes == len(file_content)
|
||||
assert result.description == "Animal relationships"
|
||||
|
||||
user_dir = service.base_dir / user.id
|
||||
stored_file = user_dir / "animals.owl"
|
||||
assert stored_file.exists()
|
||||
assert stored_file.read_bytes() == file_content
|
||||
|
||||
metadata = json.loads((user_dir / "metadata.json").read_text())
|
||||
saved_metadata = metadata["animals"]
|
||||
assert saved_metadata["filename"] == "animals.owl"
|
||||
assert saved_metadata["size_bytes"] == len(file_content)
|
||||
assert saved_metadata["description"] == "Animal relationships"
|
||||
assert saved_metadata["uploaded_at"] == result.uploaded_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_multiple_ontologies(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
service = OntologyService()
|
||||
user = SimpleNamespace(id="ontology-user")
|
||||
contents = {
|
||||
"animals": b"<rdf:RDF>Animal ontology</rdf:RDF>",
|
||||
"plants": b"<rdf:RDF>Plant ontology</rdf:RDF>",
|
||||
}
|
||||
filenames = {"animals": "animals.owl", "plants": "plants.owl"}
|
||||
descriptions = {"animals": "Animal data", "plants": "Plant data"}
|
||||
files = [
|
||||
UploadFile(filename=filenames[key], file=io.BytesIO(contents[key]))
|
||||
for key in ["animals", "plants"]
|
||||
]
|
||||
|
||||
results = await service.upload_ontologies(
|
||||
["animals", "plants"], files, user, [descriptions["animals"], descriptions["plants"]]
|
||||
)
|
||||
|
||||
assert [res.ontology_key for res in results] == ["animals", "plants"]
|
||||
for res in results:
|
||||
assert res.filename == filenames[res.ontology_key]
|
||||
assert res.size_bytes == len(contents[res.ontology_key])
|
||||
assert res.description == descriptions[res.ontology_key]
|
||||
|
||||
user_dir = service.base_dir / user.id
|
||||
metadata = json.loads((user_dir / "metadata.json").read_text())
|
||||
|
||||
for key in ["animals", "plants"]:
|
||||
stored_file = user_dir / f"{key}.owl"
|
||||
assert stored_file.exists()
|
||||
assert stored_file.read_bytes() == contents[key]
|
||||
|
||||
saved_metadata = metadata[key]
|
||||
assert saved_metadata["filename"] == filenames[key]
|
||||
assert saved_metadata["size_bytes"] == len(contents[key])
|
||||
assert saved_metadata["description"] == descriptions[key]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_ontology_contents_returns_uploaded_data(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
service = OntologyService()
|
||||
user = SimpleNamespace(id="ontology-user")
|
||||
uploads = {
|
||||
"animals": b"<rdf:RDF>Animals</rdf:RDF>",
|
||||
"plants": b"<rdf:RDF>Plants</rdf:RDF>",
|
||||
}
|
||||
|
||||
for key, content in uploads.items():
|
||||
await service.upload_ontology(
|
||||
ontology_key=key,
|
||||
file=UploadFile(filename=f"{key}.owl", file=io.BytesIO(content)),
|
||||
user=user,
|
||||
)
|
||||
|
||||
contents = service.get_ontology_contents(["animals", "plants"], user)
|
||||
|
||||
assert contents == [uploads["animals"].decode(), uploads["plants"].decode()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_ontologies_returns_metadata(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(tmp_path))
|
||||
service = OntologyService()
|
||||
user = SimpleNamespace(id="ontology-user")
|
||||
|
||||
uploads = {
|
||||
"animals": {
|
||||
"content": b"<rdf:RDF>Animals</rdf:RDF>",
|
||||
"description": "Animal ontology",
|
||||
},
|
||||
"plants": {
|
||||
"content": b"<rdf:RDF>Plants</rdf:RDF>",
|
||||
"description": "Plant ontology",
|
||||
},
|
||||
}
|
||||
|
||||
for key, payload in uploads.items():
|
||||
await service.upload_ontology(
|
||||
ontology_key=key,
|
||||
file=UploadFile(filename=f"{key}.owl", file=io.BytesIO(payload["content"])),
|
||||
user=user,
|
||||
description=payload["description"],
|
||||
)
|
||||
|
||||
metadata = service.list_ontologies(user)
|
||||
|
||||
for key, payload in uploads.items():
|
||||
entry = metadata[key]
|
||||
assert entry["filename"] == f"{key}.owl"
|
||||
assert entry["size_bytes"] == len(payload["content"])
|
||||
assert entry["description"] == payload["description"]
|
||||
|
|
@ -1,28 +1,17 @@
|
|||
import pytest
|
||||
import uuid
|
||||
from fastapi.testclient import TestClient
|
||||
from unittest.mock import Mock
|
||||
from unittest.mock import patch, Mock, AsyncMock
|
||||
from types import SimpleNamespace
|
||||
import importlib
|
||||
from cognee.api.client import app
|
||||
from cognee.modules.users.methods import get_authenticated_user
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def test_client():
|
||||
# Keep a single TestClient (and event loop) for the whole module.
|
||||
# Re-creating TestClient repeatedly can break async DB connections (asyncpg loop mismatch).
|
||||
with TestClient(app) as c:
|
||||
yield c
|
||||
gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(test_client, mock_default_user):
|
||||
async def override_get_authenticated_user():
|
||||
return mock_default_user
|
||||
|
||||
app.dependency_overrides[get_authenticated_user] = override_get_authenticated_user
|
||||
yield test_client
|
||||
app.dependency_overrides.pop(get_authenticated_user, None)
|
||||
def client():
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -43,8 +32,12 @@ def mock_default_user():
|
|||
)
|
||||
|
||||
|
||||
def test_upload_ontology_success(client):
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_ontology_success(mock_get_default_user, client, mock_default_user):
|
||||
"""Test successful ontology upload"""
|
||||
import json
|
||||
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
ontology_content = (
|
||||
b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
|
||||
)
|
||||
|
|
@ -53,7 +46,7 @@ def test_upload_ontology_success(client):
|
|||
response = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
|
||||
data={"ontology_key": unique_key, "description": "Test"},
|
||||
data={"ontology_key": json.dumps([unique_key]), "description": json.dumps(["Test"])},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
|
@ -62,8 +55,10 @@ def test_upload_ontology_success(client):
|
|||
assert "uploaded_at" in data["uploaded_ontologies"][0]
|
||||
|
||||
|
||||
def test_upload_ontology_invalid_file(client):
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user):
|
||||
"""Test 400 response for non-.owl files"""
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
|
||||
response = client.post(
|
||||
"/api/v1/ontologies",
|
||||
|
|
@ -73,10 +68,14 @@ def test_upload_ontology_invalid_file(client):
|
|||
assert response.status_code == 400
|
||||
|
||||
|
||||
def test_upload_ontology_missing_data(client):
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user):
|
||||
"""Test 400 response for missing file or key"""
|
||||
import json
|
||||
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
# Missing file
|
||||
response = client.post("/api/v1/ontologies", data={"ontology_key": "test"})
|
||||
response = client.post("/api/v1/ontologies", data={"ontology_key": json.dumps(["test"])})
|
||||
assert response.status_code == 400
|
||||
|
||||
# Missing key
|
||||
|
|
@ -86,25 +85,34 @@ def test_upload_ontology_missing_data(client):
|
|||
assert response.status_code == 400
|
||||
|
||||
|
||||
def test_upload_ontology_without_auth_header(client):
|
||||
"""Test behavior when no explicit authentication header is provided."""
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user):
|
||||
"""Test behavior when default user is provided (no explicit authentication)"""
|
||||
import json
|
||||
|
||||
unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
response = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[("ontology_file", ("test.owl", b"<rdf></rdf>", "application/xml"))],
|
||||
data={"ontology_key": unique_key},
|
||||
data={"ontology_key": json.dumps([unique_key])},
|
||||
)
|
||||
|
||||
# The current system provides a default user when no explicit authentication is given
|
||||
# This test verifies the system works with conditional authentication
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
|
||||
assert "uploaded_at" in data["uploaded_ontologies"][0]
|
||||
|
||||
|
||||
def test_upload_multiple_ontologies_in_single_request_is_rejected(client):
|
||||
"""Uploading multiple ontology files in a single request should fail."""
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_multiple_ontologies(mock_get_default_user, client, mock_default_user):
|
||||
"""Test uploading multiple ontology files in single request"""
|
||||
import io
|
||||
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
# Create mock files
|
||||
file1_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
|
||||
file2_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
|
||||
|
||||
|
|
@ -112,34 +120,45 @@ def test_upload_multiple_ontologies_in_single_request_is_rejected(client):
|
|||
("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
|
||||
("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
|
||||
]
|
||||
data = {"ontology_key": "vehicles", "description": "Base vehicles"}
|
||||
data = {
|
||||
"ontology_key": '["vehicles", "manufacturers"]',
|
||||
"descriptions": '["Base vehicles", "Car manufacturers"]',
|
||||
}
|
||||
|
||||
response = client.post("/api/v1/ontologies", files=files, data=data)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "Only one ontology_file is allowed" in response.json()["error"]
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "uploaded_ontologies" in result
|
||||
assert len(result["uploaded_ontologies"]) == 2
|
||||
assert result["uploaded_ontologies"][0]["ontology_key"] == "vehicles"
|
||||
assert result["uploaded_ontologies"][1]["ontology_key"] == "manufacturers"
|
||||
|
||||
|
||||
def test_upload_endpoint_rejects_array_style_fields(client):
|
||||
"""Array-style form values should be rejected (no backwards compatibility)."""
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_upload_endpoint_accepts_arrays(mock_get_default_user, client, mock_default_user):
|
||||
"""Test that upload endpoint accepts array parameters"""
|
||||
import io
|
||||
import json
|
||||
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
file_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
|
||||
|
||||
files = [("ontology_file", ("single.owl", io.BytesIO(file_content), "application/xml"))]
|
||||
data = {
|
||||
"ontology_key": json.dumps(["single_key"]),
|
||||
"description": json.dumps(["Single ontology"]),
|
||||
"descriptions": json.dumps(["Single ontology"]),
|
||||
}
|
||||
|
||||
response = client.post("/api/v1/ontologies", files=files, data=data)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "ontology_key must be a string" in response.json()["error"]
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["uploaded_ontologies"][0]["ontology_key"] == "single_key"
|
||||
|
||||
|
||||
def test_cognify_with_multiple_ontologies(client):
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_cognify_with_multiple_ontologies(mock_get_default_user, client, mock_default_user):
|
||||
"""Test cognify endpoint accepts multiple ontology keys"""
|
||||
payload = {
|
||||
"datasets": ["test_dataset"],
|
||||
|
|
@ -153,11 +172,14 @@ def test_cognify_with_multiple_ontologies(client):
|
|||
assert response.status_code in [200, 400, 409] # May fail for other reasons, not type
|
||||
|
||||
|
||||
def test_complete_multifile_workflow(client):
|
||||
"""Test workflow: upload ontologies one-by-one → cognify with multiple keys"""
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_complete_multifile_workflow(mock_get_default_user, client, mock_default_user):
|
||||
"""Test complete workflow: upload multiple ontologies → cognify with multiple keys"""
|
||||
import io
|
||||
import json
|
||||
|
||||
# Step 1: Upload two ontologies (one-by-one)
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
# Step 1: Upload multiple ontologies
|
||||
file1_content = b"""<?xml version="1.0"?>
|
||||
<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
|
||||
xmlns:owl="http://www.w3.org/2002/07/owl#">
|
||||
|
|
@ -170,21 +192,17 @@ def test_complete_multifile_workflow(client):
|
|||
<owl:Class rdf:ID="Manufacturer"/>
|
||||
</rdf:RDF>"""
|
||||
|
||||
upload_response_1 = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml"))],
|
||||
data={"ontology_key": "vehicles", "description": "Vehicle ontology"},
|
||||
)
|
||||
assert upload_response_1.status_code == 200
|
||||
files = [
|
||||
("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
|
||||
("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
|
||||
]
|
||||
data = {
|
||||
"ontology_key": json.dumps(["vehicles", "manufacturers"]),
|
||||
"descriptions": json.dumps(["Vehicle ontology", "Manufacturer ontology"]),
|
||||
}
|
||||
|
||||
upload_response_2 = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[
|
||||
("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml"))
|
||||
],
|
||||
data={"ontology_key": "manufacturers", "description": "Manufacturer ontology"},
|
||||
)
|
||||
assert upload_response_2.status_code == 200
|
||||
upload_response = client.post("/api/v1/ontologies", files=files, data=data)
|
||||
assert upload_response.status_code == 200
|
||||
|
||||
# Step 2: Verify ontologies are listed
|
||||
list_response = client.get("/api/v1/ontologies")
|
||||
|
|
@ -205,42 +223,44 @@ def test_complete_multifile_workflow(client):
|
|||
assert cognify_response.status_code != 400 # Not a validation error
|
||||
|
||||
|
||||
def test_upload_error_handling(client):
|
||||
"""Test error handling for invalid uploads (single-file endpoint)."""
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_multifile_error_handling(mock_get_default_user, client, mock_default_user):
|
||||
"""Test error handling for invalid multifile uploads"""
|
||||
import io
|
||||
import json
|
||||
|
||||
# Array-style key should be rejected
|
||||
# Test mismatched array lengths
|
||||
file_content = b"<rdf:RDF></rdf:RDF>"
|
||||
files = [("ontology_file", ("test.owl", io.BytesIO(file_content), "application/xml"))]
|
||||
data = {
|
||||
"ontology_key": json.dumps(["key1", "key2"]),
|
||||
"description": "desc1",
|
||||
"ontology_key": json.dumps(["key1", "key2"]), # 2 keys, 1 file
|
||||
"descriptions": json.dumps(["desc1"]),
|
||||
}
|
||||
|
||||
response = client.post("/api/v1/ontologies", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
assert "ontology_key must be a string" in response.json()["error"]
|
||||
assert "Number of keys must match number of files" in response.json()["error"]
|
||||
|
||||
# Duplicate key should be rejected
|
||||
response_1 = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml"))],
|
||||
data={"ontology_key": "duplicate", "description": "desc1"},
|
||||
)
|
||||
assert response_1.status_code == 200
|
||||
# Test duplicate keys
|
||||
files = [
|
||||
("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml")),
|
||||
("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml")),
|
||||
]
|
||||
data = {
|
||||
"ontology_key": json.dumps(["duplicate", "duplicate"]),
|
||||
"descriptions": json.dumps(["desc1", "desc2"]),
|
||||
}
|
||||
|
||||
response_2 = client.post(
|
||||
"/api/v1/ontologies",
|
||||
files=[("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml"))],
|
||||
data={"ontology_key": "duplicate", "description": "desc2"},
|
||||
)
|
||||
assert response_2.status_code == 400
|
||||
assert "already exists" in response_2.json()["error"]
|
||||
response = client.post("/api/v1/ontologies", files=files, data=data)
|
||||
assert response.status_code == 400
|
||||
assert "Duplicate ontology keys not allowed" in response.json()["error"]
|
||||
|
||||
|
||||
def test_cognify_missing_ontology_key(client):
|
||||
@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
|
||||
def test_cognify_missing_ontology_key(mock_get_default_user, client, mock_default_user):
|
||||
"""Test cognify with non-existent ontology key"""
|
||||
mock_get_default_user.return_value = mock_default_user
|
||||
|
||||
payload = {
|
||||
"datasets": ["test_dataset"],
|
||||
"ontology_key": ["nonexistent_key"],
|
||||
|
|
|
|||
|
|
@ -1,100 +0,0 @@
|
|||
import types
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from cognee.modules.search.types import SearchType
|
||||
|
||||
|
||||
def _make_user(user_id: str = "u1", tenant_id=None):
|
||||
return types.SimpleNamespace(id=user_id, tenant_id=tenant_id)
|
||||
|
||||
|
||||
def _make_dataset(*, name="ds", tenant_id="t1", dataset_id=None, owner_id=None):
|
||||
return types.SimpleNamespace(
|
||||
id=dataset_id or uuid4(),
|
||||
name=name,
|
||||
tenant_id=tenant_id,
|
||||
owner_id=owner_id or uuid4(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def search_mod():
|
||||
import importlib
|
||||
|
||||
return importlib.import_module("cognee.modules.search.methods.search")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _patch_side_effect_boundaries(monkeypatch, search_mod):
|
||||
"""
|
||||
Keep production logic; patch only unavoidable side-effect boundaries.
|
||||
"""
|
||||
|
||||
async def dummy_log_query(_query_text, _query_type, _user_id):
|
||||
return types.SimpleNamespace(id="qid-1")
|
||||
|
||||
async def dummy_log_result(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
async def dummy_prepare_search_result(search_result):
|
||||
if isinstance(search_result, tuple) and len(search_result) == 3:
|
||||
result, context, datasets = search_result
|
||||
return {"result": result, "context": context, "graphs": {}, "datasets": datasets}
|
||||
return {"result": None, "context": None, "graphs": {}, "datasets": []}
|
||||
|
||||
monkeypatch.setattr(search_mod, "send_telemetry", lambda *a, **k: None)
|
||||
monkeypatch.setattr(search_mod, "log_query", dummy_log_query)
|
||||
monkeypatch.setattr(search_mod, "log_result", dummy_log_result)
|
||||
monkeypatch.setattr(search_mod, "prepare_search_result", dummy_prepare_search_result)
|
||||
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_access_control_returns_dataset_shaped_dicts(monkeypatch, search_mod):
|
||||
user = _make_user()
|
||||
ds = _make_dataset(name="ds1", tenant_id="t1")
|
||||
|
||||
async def dummy_authorized_search(**kwargs):
|
||||
assert kwargs["dataset_ids"] == [ds.id]
|
||||
return [("r", ["ctx"], [ds])]
|
||||
|
||||
monkeypatch.setattr(search_mod, "backend_access_control_enabled", lambda: True)
|
||||
monkeypatch.setattr(search_mod, "authorized_search", dummy_authorized_search)
|
||||
|
||||
out_non_verbose = await search_mod.search(
|
||||
query_text="q",
|
||||
query_type=SearchType.CHUNKS,
|
||||
dataset_ids=[ds.id],
|
||||
user=user,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
assert out_non_verbose == [
|
||||
{
|
||||
"search_result": ["r"],
|
||||
"dataset_id": ds.id,
|
||||
"dataset_name": "ds1",
|
||||
"dataset_tenant_id": "t1",
|
||||
}
|
||||
]
|
||||
|
||||
out_verbose = await search_mod.search(
|
||||
query_text="q",
|
||||
query_type=SearchType.CHUNKS,
|
||||
dataset_ids=[ds.id],
|
||||
user=user,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
assert out_verbose == [
|
||||
{
|
||||
"search_result": ["r"],
|
||||
"dataset_id": ds.id,
|
||||
"dataset_name": "ds1",
|
||||
"dataset_tenant_id": "t1",
|
||||
"graphs": {},
|
||||
}
|
||||
]
|
||||
|
|
@ -20,30 +20,20 @@ echo "HTTP port: $HTTP_PORT"
|
|||
# smooth redeployments and container restarts while maintaining data integrity.
|
||||
echo "Running database migrations..."
|
||||
|
||||
set +e # Disable exit on error to handle specific migration errors
|
||||
MIGRATION_OUTPUT=$(alembic upgrade head)
|
||||
MIGRATION_EXIT_CODE=$?
|
||||
set -e
|
||||
|
||||
if [[ $MIGRATION_EXIT_CODE -ne 0 ]]; then
|
||||
if [[ "$MIGRATION_OUTPUT" == *"UserAlreadyExists"* ]] || [[ "$MIGRATION_OUTPUT" == *"User default_user@example.com already exists"* ]]; then
|
||||
echo "Warning: Default user already exists, continuing startup..."
|
||||
else
|
||||
echo "Migration failed with unexpected error. Trying to run Cognee without migrations."
|
||||
|
||||
echo "Initializing database tables..."
|
||||
python /app/cognee/modules/engine/operations/setup.py
|
||||
INIT_EXIT_CODE=$?
|
||||
|
||||
if [[ $INIT_EXIT_CODE -ne 0 ]]; then
|
||||
echo "Database initialization failed!"
|
||||
exit 1
|
||||
fi
|
||||
echo "Migration failed with unexpected error."
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "Database migrations done."
|
||||
fi
|
||||
|
||||
echo "Database migrations done."
|
||||
|
||||
echo "Starting server..."
|
||||
|
||||
# Add startup delay to ensure DB is ready
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
import asyncio
|
||||
import cognee
|
||||
import os
|
||||
|
||||
from pprint import pprint
|
||||
import os
|
||||
|
||||
# By default cognee uses OpenAI's gpt-5-mini LLM model
|
||||
# Provide your OpenAI LLM API KEY
|
||||
|
|
@ -25,13 +24,13 @@ async def cognee_demo():
|
|||
|
||||
# Query Cognee for information from provided document
|
||||
answer = await cognee.search("List me all the important characters in Alice in Wonderland.")
|
||||
pprint(answer)
|
||||
print(answer)
|
||||
|
||||
answer = await cognee.search("How did Alice end up in Wonderland?")
|
||||
pprint(answer)
|
||||
print(answer)
|
||||
|
||||
answer = await cognee.search("Tell me about Alice's personality.")
|
||||
pprint(answer)
|
||||
print(answer)
|
||||
|
||||
|
||||
# Cognee is an async library, it has to be called in an async context
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
from pprint import pprint
|
||||
|
||||
import cognee
|
||||
from cognee.api.v1.search import SearchType
|
||||
|
|
@ -188,7 +187,7 @@ async def main(enable_steps):
|
|||
search_results = await cognee.search(
|
||||
query_type=SearchType.GRAPH_COMPLETION, query_text="Who has experience in design tools?"
|
||||
)
|
||||
pprint(search_results)
|
||||
print(search_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import os
|
||||
import asyncio
|
||||
import pathlib
|
||||
from pprint import pprint
|
||||
|
||||
from cognee.shared.logging_utils import setup_logging, ERROR
|
||||
|
||||
import cognee
|
||||
|
|
@ -44,7 +42,7 @@ async def main():
|
|||
|
||||
# Display search results
|
||||
for result_text in search_results:
|
||||
pprint(result_text)
|
||||
print(result_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import asyncio
|
||||
import os
|
||||
from pprint import pprint
|
||||
|
||||
import cognee
|
||||
from cognee.api.v1.search import SearchType
|
||||
|
|
@ -78,7 +77,7 @@ async def main():
|
|||
query_type=SearchType.GRAPH_COMPLETION,
|
||||
query_text="What are the exact cars and their types produced by Audi?",
|
||||
)
|
||||
pprint(search_results)
|
||||
print(search_results)
|
||||
|
||||
await visualize_graph()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
import os
|
||||
import cognee
|
||||
import pathlib
|
||||
from pprint import pprint
|
||||
|
||||
from cognee.modules.users.exceptions import PermissionDeniedError
|
||||
from cognee.modules.users.tenants.methods import select_tenant
|
||||
|
|
@ -87,7 +86,7 @@ async def main():
|
|||
)
|
||||
print("\nSearch results as user_1 on dataset owned by user_1:")
|
||||
for result in search_results:
|
||||
pprint(result)
|
||||
print(f"{result}\n")
|
||||
|
||||
# But user_1 cant read the dataset owned by user_2 (QUANTUM dataset)
|
||||
print("\nSearch result as user_1 on the dataset owned by user_2:")
|
||||
|
|
@ -135,7 +134,7 @@ async def main():
|
|||
dataset_ids=[quantum_dataset_id],
|
||||
)
|
||||
for result in search_results:
|
||||
pprint(result)
|
||||
print(f"{result}\n")
|
||||
|
||||
# If we'd like for user_1 to add new documents to the QUANTUM dataset owned by user_2, user_1 would have to get
|
||||
# "write" access permission, which user_1 currently does not have
|
||||
|
|
@ -218,7 +217,7 @@ async def main():
|
|||
dataset_ids=[quantum_cognee_lab_dataset_id],
|
||||
)
|
||||
for result in search_results:
|
||||
pprint(result)
|
||||
print(f"{result}\n")
|
||||
|
||||
# Note: All of these function calls and permission system is available through our backend endpoints as well
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,4 @@
|
|||
import asyncio
|
||||
from pprint import pprint
|
||||
|
||||
import cognee
|
||||
from cognee.modules.engine.operations.setup import setup
|
||||
from cognee.modules.users.methods import get_default_user
|
||||
|
|
@ -73,7 +71,7 @@ async def main():
|
|||
print("Search results:")
|
||||
# Display results
|
||||
for result_text in search_results:
|
||||
pprint(result_text)
|
||||
print(result_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,6 +1,4 @@
|
|||
import asyncio
|
||||
from pprint import pprint
|
||||
|
||||
import cognee
|
||||
from cognee.shared.logging_utils import setup_logging, ERROR
|
||||
from cognee.api.v1.search import SearchType
|
||||
|
|
@ -56,7 +54,7 @@ async def main():
|
|||
print("Search results:")
|
||||
# Display results
|
||||
for result_text in search_results:
|
||||
pprint(result_text)
|
||||
print(result_text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
from pprint import pprint
|
||||
import cognee
|
||||
from cognee.shared.logging_utils import setup_logging, INFO
|
||||
from cognee.api.v1.search import SearchType
|
||||
|
|
@ -36,16 +35,16 @@ biography_1 = """
|
|||
|
||||
biography_2 = """
|
||||
Arnulf Øverland Ole Peter Arnulf Øverland ( 27 April 1889 – 25 March 1968 ) was a Norwegian poet and artist . He is principally known for his poetry which served to inspire the Norwegian resistance movement during the German occupation of Norway during World War II .
|
||||
|
||||
|
||||
Biography .
|
||||
Øverland was born in Kristiansund and raised in Bergen . His parents were Peter Anton Øverland ( 1852–1906 ) and Hanna Hage ( 1854–1939 ) . The early death of his father , left the family economically stressed . He was able to attend Bergen Cathedral School and in 1904 Kristiania Cathedral School . He graduated in 1907 and for a time studied philology at University of Kristiania . Øverland published his first collection of poems ( 1911 ) .
|
||||
|
||||
|
||||
Øverland became a communist sympathizer from the early 1920s and became a member of Mot Dag . He also served as chairman of the Norwegian Students Society 1923–28 . He changed his stand in 1937 , partly as an expression of dissent against the ongoing Moscow Trials . He was an avid opponent of Nazism and in 1936 he wrote the poem Du må ikke sove which was printed in the journal Samtiden . It ends with . ( I thought: : Something is imminent . Our era is over – Europe’s on fire! ) . Probably the most famous line of the poem is ( You mustnt endure so well the injustice that doesnt affect you yourself! )
|
||||
|
||||
|
||||
During the German occupation of Norway from 1940 in World War II , he wrote to inspire the Norwegian resistance movement . He wrote a series of poems which were clandestinely distributed , leading to the arrest of both him and his future wife Margrete Aamot Øverland in 1941 . Arnulf Øverland was held first in the prison camp of Grini before being transferred to Sachsenhausen concentration camp in Germany . He spent a four-year imprisonment until the liberation of Norway in 1945 . His poems were later collected in Vi overlever alt and published in 1945 .
|
||||
|
||||
|
||||
Øverland played an important role in the Norwegian language struggle in the post-war era . He became a noted supporter for the conservative written form of Norwegian called Riksmål , he was president of Riksmålsforbundet ( an organization in support of Riksmål ) from 1947 to 1956 . In addition , Øverland adhered to the traditionalist style of writing , criticising modernist poetry on several occasions . His speech Tungetale fra parnasset , published in Arbeiderbladet in 1954 , initiated the so-called Glossolalia debate .
|
||||
|
||||
|
||||
Personal life .
|
||||
In 1918 he had married the singer Hildur Arntzen ( 1888–1957 ) . Their marriage was dissolved in 1939 . In 1940 , he married Bartholine Eufemia Leganger ( 1903–1995 ) . They separated shortly after , and were officially divorced in 1945 . Øverland was married to journalist Margrete Aamot Øverland ( 1913–1978 ) during June 1945 . In 1946 , the Norwegian Parliament arranged for Arnulf and Margrete Aamot Øverland to reside at the Grotten . He lived there until his death in 1968 and she lived there for another ten years until her death in 1978 . Arnulf Øverland was buried at Vår Frelsers Gravlund in Oslo . Joseph Grimeland designed the bust of Arnulf Øverland ( bronze , 1970 ) at his grave site .
|
||||
|
||||
|
|
@ -57,7 +56,7 @@ biography_2 = """
|
|||
- Vi overlever alt ( 1945 )
|
||||
- Sverdet bak døren ( 1956 )
|
||||
- Livets minutter ( 1965 )
|
||||
|
||||
|
||||
Awards .
|
||||
- Gyldendals Endowment ( 1935 )
|
||||
- Dobloug Prize ( 1951 )
|
||||
|
|
@ -88,8 +87,7 @@ async def main():
|
|||
top_k=15,
|
||||
)
|
||||
print(f"Query: {query_text}")
|
||||
print("Results:")
|
||||
pprint(search_results)
|
||||
print(f"Results: {search_results}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import asyncio
|
||||
from pprint import pprint
|
||||
|
||||
import cognee
|
||||
from cognee.memify_pipelines.create_triplet_embeddings import create_triplet_embeddings
|
||||
|
|
@ -66,7 +65,7 @@ async def main():
|
|||
query_type=SearchType.TRIPLET_COMPLETION,
|
||||
query_text="What are the models produced by Volkswagen based on the context?",
|
||||
)
|
||||
pprint(search_results)
|
||||
print(search_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
[project]
|
||||
name = "cognee"
|
||||
|
||||
version = "0.5.1"
|
||||
version = "0.5.0.dev0"
|
||||
description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
|
||||
authors = [
|
||||
{ name = "Vasilije Markovic" },
|
||||
|
|
|
|||
4
uv.lock
generated
4
uv.lock
generated
|
|
@ -1,5 +1,5 @@
|
|||
version = 1
|
||||
revision = 3
|
||||
revision = 2
|
||||
requires-python = ">=3.10, <3.14"
|
||||
resolution-markers = [
|
||||
"python_full_version >= '3.13' and platform_python_implementation != 'PyPy' and sys_platform == 'darwin'",
|
||||
|
|
@ -946,7 +946,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "cognee"
|
||||
version = "0.5.1"
|
||||
version = "0.5.0.dev0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiofiles" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue