diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 1b6dc177..c182578a 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,17 +5,17 @@ on: workflow_dispatch: push: tags: - - v0.18.* + - v* permissions: contents: read jobs: - deploy: + python-packages: - name: Build everything - runs-on: ubuntu-latest + name: Release Python packages + runs-on: ubuntu-24.04 permissions: contents: write id-token: write @@ -25,50 +25,73 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v3 - - - name: Log in to Docker Hub - uses: docker/login-action@f4ef78c080cd8ba55a85445d5b36e214a81df20a - with: - username: ${{ vars.DOCKER_USERNAME }} - password: ${{ secrets.DOCKER_SECRET }} - - - name: Install build dependencies - run: pip3 install jsonnet + uses: actions/checkout@v4 - name: Get version id: version run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT - - run: echo ${{ steps.version.outputs.VERSION }} - - name: Build packages run: make packages VERSION=${{ steps.version.outputs.VERSION }} - name: Publish release distributions to PyPI uses: pypa/gh-action-pypi-publish@release/v1 - - name: Extract metadata for container - id: meta - uses: docker/metadata-action@v4 - with: - images: trustgraph/trustgraph-flow - tags: | - type=ref,event=branch - type=ref,event=pr - type=semver,pattern={{version}} - type=semver,pattern={{major}}.{{minor}} - type=sha + deploy-container-image: - - name: Build and push Docker image - id: push - uses: docker/build-push-action@3b5e8027fcad23fda98b2e3ac259d8d67585f671 - with: - context: . - file: ./Containerfile - push: true - tags: ${{ steps.meta.outputs.tags }} - labels: ${{ steps.meta.outputs.labels }} + name: Release container image + runs-on: ubuntu-24.04 + permissions: + contents: write + id-token: write + environment: + name: release + + steps: + + - name: Checkout + uses: actions/checkout@v4 + + - name: Docker Hub token + run: echo ${{ secrets.DOCKER_SECRET }} > docker-token.txt + + - name: Authenticate with Docker hub + run: make docker-hub-login + + - name: Get version + id: version + run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT + + - name: Put version into package manifests + run: make update-package-versions VERSION=${{ steps.version.outputs.VERSION }} + + - name: Build containers + run: make container VERSION=${{ steps.version.outputs.VERSION }} + + - name: Push containers + run: make push VERSION=${{ steps.version.outputs.VERSION }} + + release-bundle: + + name: Upload release bundle + runs-on: ubuntu-24.04 + permissions: + contents: write + id-token: write + environment: + name: release + + steps: + + - name: Checkout + uses: actions/checkout@v4 + + - name: Install build dependencies + run: pip3 install jsonnet + + - name: Get version + id: version + run: echo VERSION=$(git describe --exact-match --tags | sed 's/^v//') >> $GITHUB_OUTPUT - name: Create deploy bundle run: templates/generate-all deploy.zip ${{ steps.version.outputs.VERSION }} diff --git a/Containerfile b/Containerfile index 73c9285f..ee4e9fff 100644 --- a/Containerfile +++ b/Containerfile @@ -11,15 +11,26 @@ ENV PIP_BREAK_SYSTEM_PACKAGES=1 RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ python3-rdflib -RUN pip3 install torch --index-url https://download.pytorch.org/whl/cpu +RUN pip3 install torch==2.5.1+cpu \ + --index-url https://download.pytorch.org/whl/cpu -RUN pip3 install anthropic boto3 cohere openai google-cloud-aiplatform ollama google-generativeai \ - langchain langchain-core langchain-huggingface langchain-text-splitters \ - langchain-community pymilvus sentence-transformers transformers \ - huggingface-hub pulsar-client cassandra-driver pyyaml \ +RUN pip3 install \ + anthropic boto3 cohere openai google-cloud-aiplatform \ + ollama google-generativeai \ + langchain==0.3.13 langchain-core==0.3.28 langchain-huggingface==0.1.2 \ + langchain-text-splitters==0.3.4 \ + langchain-community==0.3.13 \ + sentence-transformers==3.4.0 transformers==4.47.1 \ + huggingface-hub==0.27.0 \ + pymilvus \ + pulsar-client==3.5.0 cassandra-driver pyyaml \ neo4j tiktoken falkordb && \ pip3 cache purge +# Most commonly used embeddings model, just build it into the container +# image +RUN huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 + # ---------------------------------------------------------------------------- # Build a container which contains the built Python packages. The build # creates a bunch of left-over cruft, a separate phase means this is only diff --git a/Makefile b/Makefile index 67094a90..80cdefc7 100644 --- a/Makefile +++ b/Makefile @@ -30,7 +30,7 @@ packages: update-package-versions pypi-upload: twine upload dist/*-${VERSION}.* -CONTAINER=docker.io/trustgraph/trustgraph-flow +CONTAINER_BASE=docker.io/trustgraph update-package-versions: mkdir -p trustgraph-cli/trustgraph @@ -44,11 +44,23 @@ update-package-versions: echo __version__ = \"${VERSION}\" > trustgraph/trustgraph/trustgraph_version.py container: update-package-versions - ${DOCKER} build -f Containerfile -t ${CONTAINER}:${VERSION} \ - --format docker + ${DOCKER} build -f containers/Containerfile.base \ + -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} . + ${DOCKER} build -f containers/Containerfile.flow \ + -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} . + ${DOCKER} build -f containers/Containerfile.bedrock \ + -t ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} . + ${DOCKER} build -f containers/Containerfile.vertexai \ + -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} . + ${DOCKER} build -f containers/Containerfile.hf \ + -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} . push: - ${DOCKER} push ${CONTAINER}:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-base:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-flow:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-bedrock:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} + ${DOCKER} push ${CONTAINER_BASE}/trustgraph-hf:${VERSION} clean: rm -rf wheels/ @@ -104,5 +116,5 @@ update-dcs: set-version docker-hub-login: cat docker-token.txt | \ - docker login -u trustgraph --password-stdin registry-1.docker.io + ${DOCKER} login -u trustgraph --password-stdin registry-1.docker.io diff --git a/containers/Containerfile.base b/containers/Containerfile.base new file mode 100644 index 00000000..b4f5bbbf --- /dev/null +++ b/containers/Containerfile.base @@ -0,0 +1,48 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp && \ + dnf clean all + +RUN pip3 install --no-cache-dir pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_cli-* && \ + rm -rf /root/wheels + +WORKDIR / + + + diff --git a/containers/Containerfile.bedrock b/containers/Containerfile.bedrock new file mode 100644 index 00000000..21819973 --- /dev/null +++ b/containers/Containerfile.bedrock @@ -0,0 +1,48 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir boto3 pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-bedrock/ /root/build/trustgraph-bedrock/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-bedrock/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_bedrock-* && \ + rm -rf /root/wheels + +WORKDIR / + diff --git a/containers/Containerfile.flow b/containers/Containerfile.flow new file mode 100644 index 00000000..8d47effe --- /dev/null +++ b/containers/Containerfile.flow @@ -0,0 +1,60 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir \ + anthropic cohere openai google-generativeai \ + ollama \ + langchain==0.3.13 langchain-core==0.3.28 \ + langchain-text-splitters==0.3.4 \ + langchain-community==0.3.13 \ + pymilvus \ + pulsar-client==3.5.0 cassandra-driver pyyaml \ + neo4j tiktoken falkordb && \ + pip3 cache purge + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-flow/ /root/build/trustgraph-flow/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-flow/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_flow-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_cli-* && \ + rm -rf /root/wheels + +WORKDIR / + diff --git a/containers/Containerfile.hf b/containers/Containerfile.hf new file mode 100644 index 00000000..4076db28 --- /dev/null +++ b/containers/Containerfile.hf @@ -0,0 +1,75 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS ai + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install torch==2.5.1+cpu \ + --index-url https://download.pytorch.org/whl/cpu + +RUN pip3 install --no-cache-dir \ + langchain==0.3.13 langchain-core==0.3.28 langchain-huggingface==0.1.2 \ + langchain-community==0.3.13 \ + sentence-transformers==3.4.0 transformers==4.47.1 \ + huggingface-hub==0.27.0 \ + pulsar-client==3.5.0 + +# Most commonly used embeddings model, just build it into the container +# image +RUN huggingface-cli download sentence-transformers/all-MiniLM-L6-v2 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM ai AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-flow/ /root/build/trustgraph-flow/ +COPY trustgraph-vertexai/ /root/build/trustgraph-vertexai/ +COPY trustgraph-bedrock/ /root/build/trustgraph-bedrock/ +COPY trustgraph-embeddings-hf/ /root/build/trustgraph-embeddings-hf/ +COPY trustgraph-cli/ /root/build/trustgraph-cli/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-flow/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-vertexai/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-bedrock/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-embeddings-hf/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-cli/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM ai + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install /root/wheels/trustgraph_base-* && \ + pip3 install /root/wheels/trustgraph_flow-* && \ + pip3 install /root/wheels/trustgraph_vertexai-* && \ + pip3 install /root/wheels/trustgraph_bedrock-* && \ + pip3 install /root/wheels/trustgraph_embeddings_hf-* && \ + pip3 install /root/wheels/trustgraph_cli-* && \ + pip3 cache purge && \ + rm -rf /root/wheels + +WORKDIR / + +CMD sleep 1000000 + diff --git a/containers/Containerfile.vertexai b/containers/Containerfile.vertexai new file mode 100644 index 00000000..72d21bde --- /dev/null +++ b/containers/Containerfile.vertexai @@ -0,0 +1,51 @@ + +# ---------------------------------------------------------------------------- +# Build an AI container. This does the torch install which is huge, and I +# like to avoid re-doing this. +# ---------------------------------------------------------------------------- + +FROM docker.io/fedora:40 AS base + +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +RUN dnf install -y python3 python3-pip python3-wheel python3-aiohttp \ + python3-rdflib + +RUN pip3 install --no-cache-dir \ + google-cloud-aiplatform pulsar-client==3.5.0 + +# ---------------------------------------------------------------------------- +# Build a container which contains the built Python packages. The build +# creates a bunch of left-over cruft, a separate phase means this is only +# needed to support package build +# ---------------------------------------------------------------------------- + +FROM base AS build + +COPY trustgraph-base/ /root/build/trustgraph-base/ +COPY trustgraph-vertexai/ /root/build/trustgraph-vertexai/ + +WORKDIR /root/build/ + +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-base/ +RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-vertexai/ + +RUN ls /root/wheels + +# ---------------------------------------------------------------------------- +# Finally, the target container. Start with base and add the package. +# ---------------------------------------------------------------------------- + +FROM base + +COPY --from=build /root/wheels /root/wheels + +RUN \ + pip3 install --no-cache-dir /root/wheels/trustgraph_base-* && \ + pip3 install --no-cache-dir /root/wheels/trustgraph_vertexai-* && \ + rm -rf /root/wheels + +WORKDIR / + + + diff --git a/docs/README.quickstart-docker-compose.md b/docs/README.quickstart-docker-compose.md index 12cb8cf0..76f7e1f5 100644 --- a/docs/README.quickstart-docker-compose.md +++ b/docs/README.quickstart-docker-compose.md @@ -113,8 +113,9 @@ Choose one of the `Docker Compose` files that meets your preferred model and gra ### AWS Bedrock API ``` -export AWS_ID_KEY= -export AWS_SECRET_KEY= +export AWS_ACCESS_KEY_ID= +export AWS_SECRET_ACCESS_KEY= +export AWS_DEFAULT_REGION= docker compose -f tg-launch-bedrock-cassandra.yaml up -d # Using Cassandra as the graph store docker compose -f tg-launch-bedrock-neo4j.yaml up -d # Using Neo4j as the graph store ``` diff --git a/grafana/dashboards/dashboard.json b/grafana/dashboards/dashboard.json index 04561863..c484dffa 100644 --- a/grafana/dashboards/dashboard.json +++ b/grafana/dashboards/dashboard.json @@ -577,7 +577,7 @@ "disableTextWrap": false, "editorMode": "builder", "exemplar": false, - "expr": "increase(processing_count_total{status!=\"success\"}[$__rate_interval])", + "expr": "sum by(job) (increase(rate_limit_count_total[$__rate_interval]))", "format": "time_series", "fullMetaSearch": false, "includeNullMetadata": true, @@ -588,7 +588,7 @@ "useBackend": false } ], - "title": "Errors", + "title": "Rate limit events", "type": "timeseries" }, { diff --git a/prometheus/prometheus.yml b/prometheus/prometheus.yml index 24102a23..0fa70314 100644 --- a/prometheus/prometheus.yml +++ b/prometheus/prometheus.yml @@ -20,6 +20,18 @@ scrape_configs: - targets: - 'pulsar:8080' + - job_name: 'bookie' + scrape_interval: 5s + static_configs: + - targets: + - 'bookie:8000' + + - job_name: 'zookeeper' + scrape_interval: 5s + static_configs: + - targets: + - 'zookeeper:8000' + - job_name: 'pdf-decoder' scrape_interval: 5s static_configs: @@ -32,11 +44,17 @@ scrape_configs: - targets: - 'chunker:8000' - - job_name: 'vectorize' + - job_name: 'document-embeddings' scrape_interval: 5s static_configs: - targets: - - 'vectorize:8000' + - 'document-embeddings:8000' + + - job_name: 'graph-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'graph-embeddings:8000' - job_name: 'embeddings' scrape_interval: 5s @@ -74,6 +92,12 @@ scrape_configs: - targets: - 'metering-rag:8000' + - job_name: 'store-doc-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'store-doc-embeddings:8000' + - job_name: 'store-graph-embeddings' scrape_interval: 5s static_configs: @@ -104,6 +128,12 @@ scrape_configs: - targets: - 'graph-rag:8000' + - job_name: 'document-rag' + scrape_interval: 5s + static_configs: + - targets: + - 'document-rag:8000' + - job_name: 'prompt' scrape_interval: 5s static_configs: @@ -122,6 +152,12 @@ scrape_configs: - targets: - 'query-graph-embeddings:8000' + - job_name: 'query-doc-embeddings' + scrape_interval: 5s + static_configs: + - targets: + - 'query-doc-embeddings:8000' + - job_name: 'query-triples' scrape_interval: 5s static_configs: @@ -145,3 +181,7 @@ scrape_configs: static_configs: - targets: - 'workbench-ui:8000' + +# Cassandra +# qdrant + diff --git a/templates/components.jsonnet b/templates/components.jsonnet index b14665d6..653c48bb 100644 --- a/templates/components.jsonnet +++ b/templates/components.jsonnet @@ -1,47 +1,70 @@ { + + // Essentials + "trustgraph-base": import "components/trustgraph.jsonnet", + "pulsar": import "components/pulsar.jsonnet", + + // LLMs "azure": import "components/azure.jsonnet", "azure-openai": import "components/azure-openai.jsonnet", "bedrock": import "components/bedrock.jsonnet", "claude": import "components/claude.jsonnet", "cohere": import "components/cohere.jsonnet", - "document-rag": import "components/document-rag.jsonnet", - "embeddings-hf": import "components/embeddings-hf.jsonnet", - "embeddings-ollama": import "components/embeddings-ollama.jsonnet", "googleaistudio": import "components/googleaistudio.jsonnet", - "grafana": import "components/grafana.jsonnet", + "ollama": import "components/ollama.jsonnet", + "openai": import "components/openai.jsonnet", + "vertexai": import "components/vertexai.jsonnet", + + // LLMs for RAG + "azure-rag": import "components/azure-rag.jsonnet", + "azure-openai-rag": import "components/azure-openai-rag.jsonnet", + "bedrock-rag": import "components/bedrock-rag.jsonnet", + "claude-rag": import "components/claude-rag.jsonnet", + "cohere-rag": import "components/cohere-rag.jsonnet", + "googleaistudio-rag": import "components/googleaistudio-rag.jsonnet", + "ollama-rag": import "components/ollama-rag.jsonnet", + "openai-rag": import "components/openai-rag.jsonnet", + "vertexai-rag": import "components/vertexai-rag.jsonnet", + + // Embeddings + "embeddings-ollama": import "components/embeddings-ollama.jsonnet", + "embeddings-hf": import "components/embeddings-hf.jsonnet", + + // Processing pipelines "graph-rag": import "components/graph-rag.jsonnet", + "document-rag": import "components/document-rag.jsonnet", + + // Vector stores + "vector-store-milvus": import "components/milvus.jsonnet", + "vector-store-qdrant": import "components/qdrant.jsonnet", + "vector-store-pinecone": import "components/pinecone.jsonnet", + + // Triples stores "triple-store-cassandra": import "components/cassandra.jsonnet", "triple-store-neo4j": import "components/neo4j.jsonnet", "triple-store-falkordb": import "components/falkordb.jsonnet", "triple-store-memgraph": import "components/memgraph.jsonnet", + + // Observability support + "grafana": import "components/grafana.jsonnet", + + // Pulsar manager is a UI for Pulsar. Uses a LOT of memory + "pulsar-manager": import "components/pulsar-manager.jsonnet", + "llamafile": import "components/llamafile.jsonnet", - "ollama": import "components/ollama.jsonnet", - "openai": import "components/openai.jsonnet", "override-recursive-chunker": import "components/chunker-recursive.jsonnet", + // The prompt manager "prompt-template": import "components/prompt-template.jsonnet", "prompt-overrides": import "components/prompt-overrides.jsonnet", - "pulsar": import "components/pulsar.jsonnet", - "pulsar-manager": import "components/pulsar-manager.jsonnet", - "trustgraph-base": import "components/trustgraph.jsonnet", - "vector-store-milvus": import "components/milvus.jsonnet", - "vector-store-qdrant": import "components/qdrant.jsonnet", - "vector-store-pinecone": import "components/pinecone.jsonnet", - "vertexai": import "components/vertexai.jsonnet", - "workbench-ui": import "components/workbench-ui.jsonnet", - "null": {}, - + // ReAct agent "agent-manager-react": import "components/agent-manager-react.jsonnet", - // FIXME: Dupes - "cassandra": import "components/cassandra.jsonnet", - "neo4j": import "components/neo4j.jsonnet", - "memgraph": import "components/memgraph.jsonnet", - "qdrant": import "components/qdrant.jsonnet", - "pinecone": import "components/pinecone.jsonnet", - "milvus": import "components/milvus.jsonnet", - "falkordb": import "components/falkordb.jsonnet", - "trustgraph": import "components/trustgraph.jsonnet", + // Optional UI + "workbench-ui": import "components/workbench-ui.jsonnet", + + // Does nothing. But, can be a hack to overwrite parameters + "null": {}, } diff --git a/templates/components/agent-manager-react.jsonnet b/templates/components/agent-manager-react.jsonnet index a995dba5..672a0439 100644 --- a/templates/components/agent-manager-react.jsonnet +++ b/templates/components/agent-manager-react.jsonnet @@ -14,11 +14,15 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("agent-manager") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "agent-manager-react", "-p", url.pulsar, + "--prompt-request-queue", + "non-persistent://tg/request/prompt-rag", + "--prompt-response-queue", + "non-persistent://tg/response/prompt-rag", "--tool-type", ] + [ tool.id + "=" + tool.type diff --git a/templates/components/azure-openai-rag.jsonnet b/templates/components/azure-openai-rag.jsonnet new file mode 100644 index 00000000..33355707 --- /dev/null +++ b/templates/components/azure-openai-rag.jsonnet @@ -0,0 +1,61 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["ollama-rag-" + key]:: value, + }, + + "azure-openai-rag-model":: "GPT-3.5-Turbo", + "azure-openai-rag-max-output-tokens":: 4192, + "azure-openai-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("azure-openai-credentials") + .with_env_var("AZURE_TOKEN", "azure-token"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-azure", + "-p", + url.pulsar, + "-x", + std.toString($["azure-openai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["azure-openai-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/azure-openai.jsonnet b/templates/components/azure-openai.jsonnet index 8afcaf11..3ecbbdac 100644 --- a/templates/components/azure-openai.jsonnet +++ b/templates/components/azure-openai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["azure-openai-" + key]:: value, + }, + "azure-openai-model":: "GPT-3.5-Turbo", "azure-openai-max-output-tokens":: 4192, "azure-openai-temperature":: 0.0, @@ -18,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure-openai", "-p", @@ -34,48 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-azure", - "-p", - url.pulsar, - "-x", - std.toString($["azure-openai-max-output-tokens"]), - "-t", - "%0.3f" % $["azure-openai-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/azure-rag.jsonnet b/templates/components/azure-rag.jsonnet new file mode 100644 index 00000000..20b7306e --- /dev/null +++ b/templates/components/azure-rag.jsonnet @@ -0,0 +1,60 @@ +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["azure-rag-" + key]:: value, + }, + + "azure-rag-max-output-tokens":: 4096, + "azure-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("azure-credentials") + .with_env_var("AZURE_TOKEN", "azure-token") + .with_env_var("AZURE_ENDPOINT", "azure-endpoint"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-azure", + "-p", + url.pulsar, + "-x", + std.toString($["azure-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["azure-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + } + +} + prompts + diff --git a/templates/components/azure.jsonnet b/templates/components/azure.jsonnet index cf10dc66..c7746e23 100644 --- a/templates/components/azure.jsonnet +++ b/templates/components/azure.jsonnet @@ -1,10 +1,14 @@ -local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["azure-" + key]:: value, + }, + "azure-max-output-tokens":: 4096, "azure-temperature":: 0.0, @@ -18,7 +22,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-azure", "-p", @@ -32,48 +36,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-azure", - "-p", - url.pulsar, - "-x", - std.toString($["azure-max-output-tokens"]), - "-t", - "%0.3f" % $["azure-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) } diff --git a/templates/components/bedrock-rag.jsonnet b/templates/components/bedrock-rag.jsonnet new file mode 100644 index 00000000..b265a9f2 --- /dev/null +++ b/templates/components/bedrock-rag.jsonnet @@ -0,0 +1,66 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; +local chunker = import "chunker-recursive.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["bedrock-rag-" + key]:: value, + }, + + "bedrock-rag-max-output-tokens":: 4096, + "bedrock-rag-temperature":: 0.0, + "bedrock-rag-model":: "mistral.mixtral-8x7b-instruct-v0:1", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("bedrock-credentials") + .with_env_var("AWS_ACCESS_KEY_ID", "aws-id-key") + .with_env_var("AWS_SECRET_ACCESS_KEY", "aws-secret") + .with_env_var("AWS_DEFAULT_REGION", "aws-region"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_bedrock) + .with_command([ + "text-completion-bedrock", + "-p", + url.pulsar, + "-x", + std.toString($["bedrock-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["bedrock-rag-temperature"], + "-m", + $["bedrock-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + chunker + diff --git a/templates/components/bedrock.jsonnet b/templates/components/bedrock.jsonnet index 6ccaa1c5..6b599057 100644 --- a/templates/components/bedrock.jsonnet +++ b/templates/components/bedrock.jsonnet @@ -6,6 +6,11 @@ local chunker = import "chunker-recursive.jsonnet"; { + with:: function(key, value) + self + { + ["bedrock-" + key]:: value, + }, + "bedrock-max-output-tokens":: 4096, "bedrock-temperature":: 0.0, "bedrock-model":: "mistral.mixtral-8x7b-instruct-v0:1", @@ -15,13 +20,13 @@ local chunker = import "chunker-recursive.jsonnet"; create:: function(engine) local envSecrets = engine.envSecrets("bedrock-credentials") - .with_env_var("AWS_ID_KEY", "aws-id-key") - .with_env_var("AWS_SECRET", "aws-secret") - .with_env_var("AWS_REGION", "aws-region"); + .with_env_var("AWS_ACCESS_KEY_ID", "aws-id-key") + .with_env_var("AWS_SECRET_ACCESS_KEY", "aws-secret") + .with_env_var("AWS_DEFAULT_REGION", "aws-region"); local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_bedrock) .with_command([ "text-completion-bedrock", "-p", @@ -37,50 +42,18 @@ local chunker = import "chunker-recursive.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-bedrock", - "-p", - url.pulsar, - "-x", - std.toString($["bedrock-max-output-tokens"]), - "-t", - "%0.3f" % $["bedrock-temperature"], - "-m", - $["bedrock-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/cassandra.jsonnet b/templates/components/cassandra.jsonnet index b52d4b04..92ecf69f 100644 --- a/templates/components/cassandra.jsonnet +++ b/templates/components/cassandra.jsonnet @@ -12,7 +12,7 @@ cassandra + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-cassandra", "-p", @@ -44,7 +44,7 @@ cassandra + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-cassandra", "-p", diff --git a/templates/components/chunker-recursive.jsonnet b/templates/components/chunker-recursive.jsonnet index 0b64b712..4a174366 100644 --- a/templates/components/chunker-recursive.jsonnet +++ b/templates/components/chunker-recursive.jsonnet @@ -14,7 +14,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("chunker") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "chunker-recursive", "-p", diff --git a/templates/components/claude-rag.jsonnet b/templates/components/claude-rag.jsonnet new file mode 100644 index 00000000..06d58db2 --- /dev/null +++ b/templates/components/claude-rag.jsonnet @@ -0,0 +1,63 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["claude-rag-" + key]:: value, + }, + + "claude-rag-model":: "claude-3-sonnet-20240229", + "claude-rag-max-output-tokens":: 4096, + "claude-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("claude-credentials") + .with_env_var("CLAUDE_KEY", "claude-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-claude", + "-p", + url.pulsar, + "-x", + std.toString($["claude-rag-max-output-tokens"]), + "-m", + $["claude-rag-model"], + "-t", + "%0.3f" % $["claude-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/claude.jsonnet b/templates/components/claude.jsonnet index 00e4ec79..e43e7504 100644 --- a/templates/components/claude.jsonnet +++ b/templates/components/claude.jsonnet @@ -5,6 +5,12 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["claude-" + key]:: value, + }, + + "claude-model":: "claude-3-sonnet-20240229", "claude-max-output-tokens":: 4096, "claude-temperature":: 0.0, @@ -13,17 +19,19 @@ local prompts = import "prompts/mixtral.jsonnet"; create:: function(engine) local envSecrets = engine.envSecrets("claude-credentials") - .with_env_var("CLAUDE_KEY_TOKEN", "claude-key"); + .with_env_var("CLAUDE_KEY", "claude-key"); local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-claude", "-p", url.pulsar, "-x", std.toString($["claude-max-output-tokens"]), + "-m", + $["claude-model"], "-t", "%0.3f" % $["claude-temperature"], ]) @@ -31,48 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-claude", - "-p", - url.pulsar, - "-x", - std.toString($["claude-max-output-tokens"]), - "-t", - "%0.3f" % $["claude-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/cohere-rag.jsonnet b/templates/components/cohere-rag.jsonnet new file mode 100644 index 00000000..6a142519 --- /dev/null +++ b/templates/components/cohere-rag.jsonnet @@ -0,0 +1,56 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["cohere-rag-" + key]:: value, + }, + + "cohere-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("cohere-credentials") + .with_env_var("COHERE_KEY", "cohere-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-cohere", + "-p", + url.pulsar, + "-t", + "%0.3f" % $["cohere-rag-temperature"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/cohere.jsonnet b/templates/components/cohere.jsonnet index 5bc9b39c..093436fd 100644 --- a/templates/components/cohere.jsonnet +++ b/templates/components/cohere.jsonnet @@ -5,9 +5,10 @@ local prompts = import "prompts/mixtral.jsonnet"; { - // Override chunking - "chunk-size":: 150, - "chunk-overlap":: 10, + with:: function(key, value) + self + { + ["cohere-" + key]:: value, + }, "cohere-temperature":: 0.0, @@ -20,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-cohere", "-p", @@ -31,45 +32,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-cohere", - "-p", - url.pulsar, - "-t", - "%0.3f" % $["cohere-temperature"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/document-rag.jsonnet b/templates/components/document-rag.jsonnet index 0a68dd52..ec125ed5 100644 --- a/templates/components/document-rag.jsonnet +++ b/templates/components/document-rag.jsonnet @@ -11,7 +11,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("document-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "document-rag", "-p", @@ -39,5 +39,35 @@ local prompts = import "prompts/mixtral.jsonnet"; }, + "document-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("document-embeddings") + .with_image(images.trustgraph_flow) + .with_command([ + "document-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "document-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/embeddings-fastembed.jsonnet b/templates/components/embeddings-fastembed.jsonnet new file mode 100644 index 00000000..c1fe35ff --- /dev/null +++ b/templates/components/embeddings-fastembed.jsonnet @@ -0,0 +1,43 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + "embeddings-model":: "sentence-transformers/all-MiniLM-L6-v2", + + embeddings +: { + + create:: function(engine) + + local container = + engine.container("embeddings") + .with_image(images.trustgraph_flow) + .with_command([ + "embeddings-fastembed", + "-p", + url.pulsar, + "-m", + $["embeddings-model"], + ]) + .with_limits("1.0", "400M") + .with_reservations("0.5", "400M"); + + local containerSet = engine.containers( + "embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + +} + diff --git a/templates/components/embeddings-hf.jsonnet b/templates/components/embeddings-hf.jsonnet index b46feac7..29ebbc48 100644 --- a/templates/components/embeddings-hf.jsonnet +++ b/templates/components/embeddings-hf.jsonnet @@ -13,7 +13,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_hf) .with_command([ "embeddings-hf", "-p", diff --git a/templates/components/embeddings-ollama.jsonnet b/templates/components/embeddings-ollama.jsonnet index 425a1c47..a26ad0ba 100644 --- a/templates/components/embeddings-ollama.jsonnet +++ b/templates/components/embeddings-ollama.jsonnet @@ -13,7 +13,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "embeddings-ollama", "-p", diff --git a/templates/components/falkordb.jsonnet b/templates/components/falkordb.jsonnet index e238cebe..c08896d3 100644 --- a/templates/components/falkordb.jsonnet +++ b/templates/components/falkordb.jsonnet @@ -13,7 +13,7 @@ falkordb + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-falkordb", "-p", @@ -45,7 +45,7 @@ falkordb + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-falkordb", "-p", diff --git a/templates/components/googleaistudio-rag.jsonnet b/templates/components/googleaistudio-rag.jsonnet new file mode 100644 index 00000000..332749e8 --- /dev/null +++ b/templates/components/googleaistudio-rag.jsonnet @@ -0,0 +1,65 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["googleaistudio-rag-" + key]:: value, + }, + + "googleaistudio-rag-max-output-tokens":: 4096, + "googleaistudio-rag-temperature":: 0.0, + "googleaistudio-rag-model":: "gemini-1.5-flash-002", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("googleaistudio-credentials") + .with_env_var("GOOGLE_AI_STUDIO_KEY", "googleaistudio-key"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-googleaistudio", + "-p", + url.pulsar, + "-x", + std.toString( + $["googleaistudio-rag-max-output-tokens"] + ), + "-t", + "%0.3f" % $["googleaistudio-rag-temperature"], + "-m", + $["googleaistudio-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/googleaistudio.jsonnet b/templates/components/googleaistudio.jsonnet index c2a40f2c..58c7807d 100644 --- a/templates/components/googleaistudio.jsonnet +++ b/templates/components/googleaistudio.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["googleaistudio-" + key]:: value, + }, + "googleaistudio-max-output-tokens":: 4096, "googleaistudio-temperature":: 0.0, "googleaistudio-model":: "gemini-1.5-flash-002", @@ -13,12 +18,12 @@ local prompts = import "prompts/mixtral.jsonnet"; create:: function(engine) - local envSecrets = engine.envSecrets("googleaistudio-key") + local envSecrets = engine.envSecrets("googleaistudio-credentials") .with_env_var("GOOGLE_AI_STUDIO_KEY", "googleaistudio-key"); local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-googleaistudio", "-p", @@ -34,50 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-googleaistudio", - "-p", - url.pulsar, - "-x", - std.toString($["googleaistudio-max-output-tokens"]), - "-t", - "%0.3f" % $["googleaistudio-temperature"], - "-m", - $["googleaistudio-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8000, 8000, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8000, 8000, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/graph-rag.jsonnet b/templates/components/graph-rag.jsonnet index 860152c9..27035b35 100644 --- a/templates/components/graph-rag.jsonnet +++ b/templates/components/graph-rag.jsonnet @@ -14,7 +14,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-definitions") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-definitions", "-p", @@ -44,7 +44,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-relationships") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-relationships", "-p", @@ -74,7 +74,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("kg-extract-topics") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "kg-extract-topics", "-p", @@ -104,7 +104,7 @@ local url = import "values/url.jsonnet"; local container = engine.container("graph-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "graph-rag", "-p", @@ -138,5 +138,35 @@ local url = import "values/url.jsonnet"; }, + "graph-embeddings" +: { + + create:: function(engine) + + local container = + engine.container("graph-embeddings") + .with_image(images.trustgraph_flow) + .with_command([ + "graph-embeddings", + "-p", + url.pulsar, + ]) + .with_limits("1.0", "512M") + .with_reservations("0.5", "512M"); + + local containerSet = engine.containers( + "graph-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + containerSet, + service, + ]) + + }, + } diff --git a/templates/components/llamafile-rag.jsonnet b/templates/components/llamafile-rag.jsonnet new file mode 100644 index 00000000..262f586e --- /dev/null +++ b/templates/components/llamafile-rag.jsonnet @@ -0,0 +1,57 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/slm.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["llamafile-rag-" + key]:: value, + }, + + "llamafile-rag-model":: "LLaMA_CPP", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("llamafile-credentials") + .with_env_var("LLAMAFILE_URL", "llamafile-url"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-llamafile", + "-p", + url.pulsar, + "-m", + $["llamafile-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/llamafile.jsonnet b/templates/components/llamafile.jsonnet index bc1a011c..f3e1efd3 100644 --- a/templates/components/llamafile.jsonnet +++ b/templates/components/llamafile.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/slm.jsonnet"; { + with:: function(key, value) + self + { + ["llamafile-" + key]:: value, + }, + "llamafile-model":: "LLaMA_CPP", "text-completion" +: { @@ -16,7 +21,7 @@ local prompts = import "prompts/slm.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-llamafile", "-p", @@ -28,46 +33,18 @@ local prompts = import "prompts/slm.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-llamafile", - "-p", - url.pulsar, - "-m", - $["llamafile-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/memgraph.jsonnet b/templates/components/memgraph.jsonnet index 609da3a2..21684a61 100644 --- a/templates/components/memgraph.jsonnet +++ b/templates/components/memgraph.jsonnet @@ -14,7 +14,7 @@ memgraph + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-memgraph", "-p", @@ -48,7 +48,7 @@ memgraph + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-memgraph", "-p", diff --git a/templates/components/milvus.jsonnet b/templates/components/milvus.jsonnet index b3044f98..27e5e316 100644 --- a/templates/components/milvus.jsonnet +++ b/templates/components/milvus.jsonnet @@ -12,7 +12,7 @@ milvus + { local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-milvus", "-p", @@ -44,7 +44,7 @@ milvus + { local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-milvus", "-p", @@ -76,7 +76,7 @@ milvus + { local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-milvus", "-p", @@ -108,7 +108,7 @@ milvus + { local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-milvus", "-p", diff --git a/templates/components/neo4j.jsonnet b/templates/components/neo4j.jsonnet index b70562fe..7cebdc71 100644 --- a/templates/components/neo4j.jsonnet +++ b/templates/components/neo4j.jsonnet @@ -13,7 +13,7 @@ neo4j + { local container = engine.container("store-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-write-neo4j", "-p", @@ -45,7 +45,7 @@ neo4j + { local container = engine.container("query-triples") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "triples-query-neo4j", "-p", diff --git a/templates/components/ollama-rag.jsonnet b/templates/components/ollama-rag.jsonnet new file mode 100644 index 00000000..680adea5 --- /dev/null +++ b/templates/components/ollama-rag.jsonnet @@ -0,0 +1,57 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["ollama-rag-" + key]:: value, + }, + + "ollama-rag-model":: "gemma2:9b", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("ollama-credentials") + .with_env_var("OLLAMA_HOST", "ollama-host"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-ollama", + "-p", + url.pulsar, + "-m", + $["ollama-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/ollama.jsonnet b/templates/components/ollama.jsonnet index 8da00848..95f1abf0 100644 --- a/templates/components/ollama.jsonnet +++ b/templates/components/ollama.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["ollama-" + key]:: value, + }, + "ollama-model":: "gemma2:9b", "text-completion" +: { @@ -16,7 +21,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-ollama", "-p", @@ -28,46 +33,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-ollama", - "-p", - url.pulsar, - "-m", - $["ollama-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/openai-rag.jsonnet b/templates/components/openai-rag.jsonnet new file mode 100644 index 00000000..bfb7dd98 --- /dev/null +++ b/templates/components/openai-rag.jsonnet @@ -0,0 +1,63 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["openai-rag-" + key]:: value, + }, + + "openai-rag-max-output-tokens":: 4096, + "openai-rag-temperature":: 0.0, + "openai-rag-model":: "GPT-3.5-Turbo", + + "text-completion-rag" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("openai-credentials") + .with_env_var("OPENAI_TOKEN", "openai-token"); + + local containerRag = + engine.container("text-completion-rag") + .with_image(images.trustgraph_flow) + .with_command([ + "text-completion-openai", + "-p", + url.pulsar, + "-x", + std.toString($["openai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["openai-rag-temperature"], + "-m", + $["openai-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSetRag = engine.containers( + "text-completion-rag", [ containerRag ] + ); + + local serviceRag = + engine.internalService(containerSetRag) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSetRag, + serviceRag, + ]) + + }, + +} + prompts + diff --git a/templates/components/openai.jsonnet b/templates/components/openai.jsonnet index 27725cb6..9e0212d2 100644 --- a/templates/components/openai.jsonnet +++ b/templates/components/openai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["openai-" + key]:: value, + }, + "openai-max-output-tokens":: 4096, "openai-temperature":: 0.0, "openai-model":: "GPT-3.5-Turbo", @@ -18,7 +23,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "text-completion-openai", "-p", @@ -34,50 +39,18 @@ local prompts = import "prompts/mixtral.jsonnet"; .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); - local containerRag = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-openai", - "-p", - url.pulsar, - "-x", - std.toString($["openai-max-output-tokens"]), - "-t", - "%0.3f" % $["openai-temperature"], - "-m", - $["openai-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_env_var_secrets(envSecrets) - .with_limits("0.5", "128M") - .with_reservations("0.1", "128M"); - local containerSet = engine.containers( "text-completion", [ container ] ); - local containerSetRag = engine.containers( - "text-completion-rag", [ containerRag ] - ); - local service = engine.internalService(containerSet) .with_port(8080, 8080, "metrics"); - local serviceRag = - engine.internalService(containerSetRag) - .with_port(8080, 8080, "metrics"); - engine.resources([ envSecrets, containerSet, - containerSetRag, service, - serviceRag, ]) }, diff --git a/templates/components/pinecone.jsonnet b/templates/components/pinecone.jsonnet index 3422952a..ede383a5 100644 --- a/templates/components/pinecone.jsonnet +++ b/templates/components/pinecone.jsonnet @@ -17,7 +17,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-pinecone", "-p", @@ -52,7 +52,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-pinecone", "-p", @@ -87,7 +87,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-pinecone", "-p", @@ -122,7 +122,7 @@ local cassandra_hosts = "cassandra"; local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-pinecone", "-p", diff --git a/templates/components/prompt-overrides.jsonnet b/templates/components/prompt-overrides.jsonnet index 648e5b66..852ec09d 100644 --- a/templates/components/prompt-overrides.jsonnet +++ b/templates/components/prompt-overrides.jsonnet @@ -1,7 +1,3 @@ -local base = import "base/base.jsonnet"; -local images = import "values/images.jsonnet"; -local url = import "values/url.jsonnet"; -local prompts = import "prompts/mixtral.jsonnet"; local default_prompts = import "prompts/default-prompts.jsonnet"; { diff --git a/templates/components/prompt-template.jsonnet b/templates/components/prompt-template.jsonnet index 3dadf337..b3187c9b 100644 --- a/templates/components/prompt-template.jsonnet +++ b/templates/components/prompt-template.jsonnet @@ -44,7 +44,7 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("prompt") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "prompt-template", "-p", @@ -84,7 +84,7 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; local container = engine.container("prompt-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "prompt-template", "-p", diff --git a/templates/components/pulsar.jsonnet b/templates/components/pulsar.jsonnet index 0342b4d5..d111f616 100644 --- a/templates/components/pulsar.jsonnet +++ b/templates/components/pulsar.jsonnet @@ -2,32 +2,114 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; +// This is a Pulsar configuration. Non-standalone mode so we deploy +// individual components: bookkeeper, broker and zookeeper. +// +// This also deploys the TrustGraph 'admin' container which initialises +// TrustGraph-specific namespaces etc. + { "pulsar" +: { create:: function(engine) -// local confVolume = engine.volume("pulsar-conf").with_size("2G"); - local dataVolume = engine.volume("pulsar-data").with_size("20G"); + // Zookeeper volume + local zkVolume = engine.volume("zookeeper").with_size("1G"); - local container = + // Zookeeper container + local zkContainer = + engine.container("zookeeper") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/zookeeper.conf && bin/generate-zookeeper-config.sh conf/zookeeper.conf && exec bin/pulsar zookeeper" + ]) + .with_limits("1", "400M") + .with_reservations("0.05", "400M") + .with_user("0:1000") + .with_volume_mount(zkVolume, "/pulsar/data/zookeeper") + .with_environment({ + "metadataStoreUrl": "zk:zookeeper:2181", + "PULSAR_MEM": "-Xms256m -Xmx256m -XX:MaxDirectMemorySize=256m", + }) + .with_port(2181, 2181, "zookeeper") + .with_port(2888, 2888, "zookeeper2") + .with_port(3888, 3888, "zookeeper3"); + + // Pulsar cluster init container + local initContainer = + engine.container("pulsar-init") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "sleep 10 && bin/pulsar initialize-cluster-metadata --cluster cluster-a --zookeeper zookeeper:2181 --configuration-store zookeeper:2181 --web-service-url http://pulsar:8080 --broker-service-url pulsar://pulsar:6650", + ]) + .with_limits("1", "512M") + .with_reservations("0.05", "512M") + .with_environment({ + "PULSAR_MEM": "-Xms256m -Xmx256m -XX:MaxDirectMemorySize=256m", + }); + + + // Bookkeeper volume + local bookieVolume = engine.volume("bookie").with_size("20G"); + + // Bookkeeper container + local bookieContainer = + engine.container("bookie") + .with_image(images.pulsar) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/bookkeeper.conf && exec bin/pulsar bookie" + // false ^ causes this to be a 'failure' exit. + ]) + .with_limits("1", "800M") + .with_reservations("0.1", "800M") + .with_user("0:1000") + .with_volume_mount(bookieVolume, "/pulsar/data/bookkeeper") + .with_environment({ + "clusterName": "cluster-a", + "zkServers": "zookeeper:2181", + "bookieId": "bookie", + "metadataStoreUri": "metadata-store:zk:zookeeper:2181", + "advertisedAddress": "bookie", + "BOOKIE_MEM": "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=256m", + }) + .with_port(3181, 3181, "bookie"); + + // Pulsar broker, stateless (uses ZK and Bookkeeper for state) + local brokerContainer = engine.container("pulsar") .with_image(images.pulsar) - .with_command(["bin/pulsar", "standalone"]) + .with_command([ + "bash", + "-c", + "bin/apply-config-from-env.py conf/broker.conf && exec bin/pulsar broker" + ]) + .with_limits("1", "800M") + .with_reservations("0.1", "800M") .with_environment({ - "PULSAR_MEM": "-Xms600M -Xmx600M" + "metadataStoreUrl": "zk:zookeeper:2181", + "zookeeperServers": "zookeeper:2181", + "clusterName": "cluster-a", + "managedLedgerDefaultEnsembleSize": "1", + "managedLedgerDefaultWriteQuorum": "1", + "managedLedgerDefaultAckQuorum": "1", + "advertisedAddress": "pulsar", + "advertisedListeners": "external:pulsar://pulsar:6650,localhost:pulsar://localhost:6650", + "PULSAR_MEM": "-Xms512m -Xmx512m -XX:MaxDirectMemorySize=256m", }) - .with_limits("2.0", "1500M") - .with_reservations("1.0", "1500M") -// .with_volume_mount(confVolume, "/pulsar/conf") - .with_volume_mount(dataVolume, "/pulsar/data") - .with_port(6650, 6650, "bookie") - .with_port(8080, 8080, "http"); + .with_port(6650, 6650, "pulsar") + .with_port(8080, 8080, "admin"); + // Trustgraph Pulsar initialisation local adminContainer = - engine.container("init-pulsar") - .with_image(images.trustgraph) + engine.container("init-trustgraph") + .with_image(images.trustgraph_flow) .with_command([ "tg-init-pulsar", "-p", @@ -36,10 +118,32 @@ local url = import "values/url.jsonnet"; .with_limits("1", "128M") .with_reservations("0.1", "128M"); - local containerSet = engine.containers( + // Container sets + local zkContainerSet = engine.containers( + "zookeeper", + [ + zkContainer, + ] + ); + + local initContainerSet = engine.containers( + "init-pulsar", + [ + initContainer, + ] + ); + + local bookieContainerSet = engine.containers( + "bookie", + [ + bookieContainer, + ] + ); + + local brokerContainerSet = engine.containers( "pulsar", [ - container + brokerContainer, ] ); @@ -50,17 +154,35 @@ local url = import "values/url.jsonnet"; ] ); - local service = - engine.service(containerSet) - .with_port(6650, 6650, "bookie") - .with_port(8080, 8080, "http"); + // Zookeeper service + local zkService = + engine.service(zkContainerSet) + .with_port(2181, 2181, "zookeeper") + .with_port(2888, 2888, "zookeeper2") + .with_port(3888, 3888, "zookeeper3"); + + // Bookkeeper service + local bookieService = + engine.service(bookieContainerSet) + .with_port(3181, 3181, "bookie"); + + // Pulsar broker service + local brokerService = + engine.service(brokerContainerSet) + .with_port(6650, 6650, "pulsar") + .with_port(8080, 8080, "admin"); engine.resources([ -// confVolume, - dataVolume, - containerSet, + zkVolume, + bookieVolume, + zkContainerSet, + initContainerSet, + bookieContainerSet, + brokerContainerSet, adminContainerSet, - service, + zkService, + bookieService, + brokerService, ]) } diff --git a/templates/components/qdrant.jsonnet b/templates/components/qdrant.jsonnet index f923e84f..352cb741 100644 --- a/templates/components/qdrant.jsonnet +++ b/templates/components/qdrant.jsonnet @@ -12,7 +12,7 @@ qdrant + { local container = engine.container("store-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-write-qdrant", "-p", @@ -44,7 +44,7 @@ qdrant + { local container = engine.container("query-graph-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "ge-query-qdrant", "-p", @@ -76,7 +76,7 @@ qdrant + { local container = engine.container("store-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-write-qdrant", "-p", @@ -108,7 +108,7 @@ qdrant + { local container = engine.container("query-doc-embeddings") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "de-query-qdrant", "-p", diff --git a/templates/components/trustgraph.jsonnet b/templates/components/trustgraph.jsonnet index 31ae420e..833d932b 100644 --- a/templates/components/trustgraph.jsonnet +++ b/templates/components/trustgraph.jsonnet @@ -1,7 +1,6 @@ local base = import "base/base.jsonnet"; local images = import "values/images.jsonnet"; local url = import "values/url.jsonnet"; -local prompt = import "prompt-template.jsonnet"; { @@ -22,7 +21,7 @@ local prompt = import "prompt-template.jsonnet"; local container = engine.container("api-gateway") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "api-gateway", "-p", @@ -61,7 +60,7 @@ local prompt = import "prompt-template.jsonnet"; local container = engine.container("chunker") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "chunker-token", "-p", @@ -95,7 +94,7 @@ local prompt = import "prompt-template.jsonnet"; local container = engine.container("pdf-decoder") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "pdf-decoder", "-p", @@ -119,43 +118,13 @@ local prompt = import "prompt-template.jsonnet"; }, - "vectorize" +: { - - create:: function(engine) - - local container = - engine.container("vectorize") - .with_image(images.trustgraph) - .with_command([ - "embeddings-vectorize", - "-p", - url.pulsar, - ]) - .with_limits("1.0", "512M") - .with_reservations("0.5", "512M"); - - local containerSet = engine.containers( - "vectorize", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - containerSet, - service, - ]) - - }, - "metering" +: { create:: function(engine) local container = engine.container("metering") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "metering", "-p", @@ -185,7 +154,7 @@ local prompt = import "prompt-template.jsonnet"; local container = engine.container("metering-rag") - .with_image(images.trustgraph) + .with_image(images.trustgraph_flow) .with_command([ "metering", "-p", @@ -211,5 +180,5 @@ local prompt = import "prompt-template.jsonnet"; }, -} + prompt +} diff --git a/templates/components/vertexai-rag.jsonnet b/templates/components/vertexai-rag.jsonnet new file mode 100644 index 00000000..0b5cf9a3 --- /dev/null +++ b/templates/components/vertexai-rag.jsonnet @@ -0,0 +1,74 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local prompts = import "prompts/mixtral.jsonnet"; + +{ + + with:: function(key, value) + self + { + ["vertexai-rag-" + key]:: value, + }, + + "vertexai-rag-model":: "gemini-1.0-pro-001", + "vertexai-rag-private-key":: "/vertexai/private.json", + "vertexai-rag-region":: "us-central1", + "vertexai-rag-max-output-tokens":: 4096, + "vertexai-rag-temperature":: 0.0, + + "text-completion-rag" +: { + + create:: function(engine) + + local cfgVol = engine.secretVolume( + "vertexai-creds", + "./vertexai", + { + "private.json": importstr "vertexai/private.json", + } + ); + + local container = + engine.container("text-completion-rag") + .with_image(images.trustgraph_vertexai) + .with_command([ + "text-completion-vertexai", + "-p", + url.pulsar, + "-k", + $["vertexai-rag-private-key"], + "-r", + $["vertexai-rag-region"], + "-x", + std.toString($["vertexai-rag-max-output-tokens"]), + "-t", + "%0.3f" % $["vertexai-rag-temperature"], + "-m", + $["vertexai-rag-model"], + "-i", + "non-persistent://tg/request/text-completion-rag", + "-o", + "non-persistent://tg/response/text-completion-rag", + ]) + .with_limits("0.5", "256M") + .with_reservations("0.1", "256M") + .with_volume_mount(cfgVol, "/vertexai"); + + local containerSet = engine.containers( + "text-completion-rag", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8000, 8000, "metrics"); + + engine.resources([ + cfgVol, + containerSet, + service, + ]) + + } + +} + prompts + diff --git a/templates/components/vertexai.jsonnet b/templates/components/vertexai.jsonnet index ef193156..0e3550c5 100644 --- a/templates/components/vertexai.jsonnet +++ b/templates/components/vertexai.jsonnet @@ -5,6 +5,11 @@ local prompts = import "prompts/mixtral.jsonnet"; { + with:: function(key, value) + self + { + ["vertexai-" + key]:: value, + }, + "vertexai-model":: "gemini-1.0-pro-001", "vertexai-private-key":: "/vertexai/private.json", "vertexai-region":: "us-central1", @@ -25,7 +30,7 @@ local prompts = import "prompts/mixtral.jsonnet"; local container = engine.container("text-completion") - .with_image(images.trustgraph) + .with_image(images.trustgraph_vertexai) .with_command([ "text-completion-vertexai", "-p", @@ -61,59 +66,5 @@ local prompts = import "prompts/mixtral.jsonnet"; }, - "text-completion-rag" +: { - - create:: function(engine) - - local cfgVol = engine.secretVolume( - "vertexai-creds", - "./vertexai", - { - "private.json": importstr "vertexai/private.json", - } - ); - - local container = - engine.container("text-completion-rag") - .with_image(images.trustgraph) - .with_command([ - "text-completion-vertexai", - "-p", - url.pulsar, - "-k", - $["vertexai-private-key"], - "-r", - $["vertexai-region"], - "-x", - std.toString($["vertexai-max-output-tokens"]), - "-t", - "%0.3f" % $["vertexai-temperature"], - "-m", - $["vertexai-model"], - "-i", - "non-persistent://tg/request/text-completion-rag", - "-o", - "non-persistent://tg/response/text-completion-rag", - ]) - .with_limits("0.5", "256M") - .with_reservations("0.1", "256M") - .with_volume_mount(cfgVol, "/vertexai"); - - local containerSet = engine.containers( - "text-completion-rag", [ container ] - ); - - local service = - engine.internalService(containerSet) - .with_port(8000, 8000, "metrics"); - - engine.resources([ - cfgVol, - containerSet, - service, - ]) - - } - } + prompts diff --git a/templates/engine/docker-compose.jsonnet b/templates/engine/docker-compose.jsonnet index c37f1df0..0be3c3e3 100644 --- a/templates/engine/docker-compose.jsonnet +++ b/templates/engine/docker-compose.jsonnet @@ -22,6 +22,8 @@ with_image:: function(x) self + { image: x }, + with_user:: function(x) self + { user: x }, + with_command:: function(x) self + { command: x }, with_environment:: function(x) self + { @@ -75,6 +77,10 @@ { command: container.command } else {}) + + (if std.objectHas(container, "user") then + { user: container.user } + else {}) + + (if ! std.isEmpty(container.environment) then { environment: container.environment } else {}) + diff --git a/templates/engine/k8s.jsonnet b/templates/engine/k8s.jsonnet index 2fec0d1f..dfd8d11f 100644 --- a/templates/engine/k8s.jsonnet +++ b/templates/engine/k8s.jsonnet @@ -14,6 +14,8 @@ with_image:: function(x) self + { image: x }, + with_user:: function(x) self + { user: x }, + with_command:: function(x) self + { command: x }, with_environment:: function(x) self + { diff --git a/templates/generate b/templates/generate index e8772483..6f915cda 100755 --- a/templates/generate +++ b/templates/generate @@ -7,6 +7,8 @@ import logging import os import sys import zipfile +import pathlib +from io import BytesIO logger = logging.getLogger("generate") logging.basicConfig(level=logging.INFO, format='%(message)s') @@ -15,9 +17,13 @@ private_json = "Put your GCP private.json here" class Generator: - def __init__(self, config, base="./templates/", version="0.0.0"): + def __init__( + self, config, templates="./templates/", resources="./resources", + version="0.0.0", + ): - self.jsonnet_base = base + self.templates = pathlib.Path(templates) + self.resources = pathlib.Path(resources) self.config = config self.version = f"\"{version}\"".encode("utf-8") @@ -34,25 +40,27 @@ class Generator: path = os.path.join(".", dir, filename) return str(path), self.config - if filename == "version.jsonnet" and dir == "./templates/values/": + if filename == "version.jsonnet" and dir == "templates/values/": path = os.path.join(".", dir, filename) return str(path), self.version if dir: candidates = [ - os.path.join(".", dir, filename), - os.path.join(".", filename) + self.templates.joinpath(dir, filename), + self.templates.joinpath(filename), + self.resources.joinpath(dir, filename), + self.resources.joinpath(filename), ] else: candidates = [ - os.path.join(".", filename) + self.templates.joinpath(filename), ] try: if filename == "vertexai/private.json": - return candidates[0], private_json.encode("utf-8") + return str(candidates[0]), private_json.encode("utf-8") for c in candidates: logger.debug("Try: %s", c) @@ -68,73 +76,157 @@ class Generator: except: - path = os.path.join(self.jsonnet_base, filename) + path = os.path.join(self.templates, filename) logger.debug("Try: %s", path) with open(path, "rb") as f: logger.debug("Loaded: %s", path) return str(path), f.read() +class Packager: + + def __init__(self): + self.templates = pathlib.Path("./templates") + self.resources = pathlib.Path("./resources") + + def process( + self, config, version="0.0.0", platform="docker-compose", + ): + + config = config.encode("utf-8") + + gen = Generator( + config, templates=self.templates, resources=self.resources, + version=version + ) + + path = self.templates.joinpath( + f"config-to-{platform}.jsonnet" + ) + wrapper = path.read_text() + + processed = gen.process(wrapper) + + return processed + + def generate(self, config, version, platform): + + logger.info(f"Generating for platform={platform} version={version}") + + try: + + if platform in set(["docker-compose", "podman-compose"]): + return self.generate_docker_compose( + "docker-compose", version, config + ) + elif platform in set(["minikube-k8s", "gcp-k8s"]): + return self.generate_k8s( + platform, version, config + ) + else: + raise RuntimeError("Bad configuration") + + except Exception as e: + logging.error(f"Exception: {e}") + raise e + + def generate_docker_compose(self, platform, version, config): + + processed = self.process( + config, platform=platform, version=version + ) + + y = yaml.dump(processed) + + mem = BytesIO() + + with zipfile.ZipFile(mem, mode='w') as out: + + def output(name, content): + logger.info(f"Adding {name}...") + out.writestr(name, content) + + fname = "docker-compose.yaml" + + output(fname, y) + + # Grafana config + path = self.resources.joinpath( + "grafana/dashboards/dashboard.json" + ) + res = path.read_text() + output("grafana/dashboards/dashboard.json", res) + + path = self.resources.joinpath( + "grafana/provisioning/dashboard.yml" + ) + res = path.read_text() + output("grafana/provisioning/dashboard.yml", res) + + path = self.resources.joinpath( + "grafana/provisioning/datasource.yml" + ) + res = path.read_text() + output("grafana/provisioning/datasource.yml", res) + + # Prometheus config + path = self.resources.joinpath( + "prometheus/prometheus.yml" + ) + res = path.read_text() + output("prometheus/prometheus.yml", res) + + logger.info("Generation complete.") + + return mem.getvalue() + + def generate_k8s(self, platform, version, config): + + processed = self.process( + config, platform=platform, version=version + ) + + y = yaml.dump(processed) + + mem = BytesIO() + + with zipfile.ZipFile(mem, mode='w') as out: + + def output(name, content): + logger.info(f"Adding {name}...") + out.writestr(name, content) + + fname = "resources.yaml" + + output(fname, y) + + logger.info("Generation complete.") + + return mem.getvalue() + def main(): - if len(sys.argv) != 3: + if len(sys.argv) != 4: print() print("Usage:") - print(" generate < input.json") + print(" generate < input.json") print() sys.exit(1) outfile = sys.argv[1] version = sys.argv[2] + platform = sys.argv[3] cfg = sys.stdin.read() - cfg = json.loads(cfg) logger.info(f"Outputting to {outfile}...") - with zipfile.ZipFile(outfile, mode='w') as out: + p = Packager() + resp = p.generate(cfg, version, platform) - def output(name, content): - logger.info(f"Adding {name}...") - out.writestr(name, content) + with open(outfile, "wb") as f: + f.write(resp) - fname = "tg-launch.yaml" - - platform = "docker-compose" - - with open(f"./templates/config-to-{platform}.jsonnet", "r") as f: - wrapper = f.read() - - gen = Generator(json.dumps(cfg).encode("utf-8"), version=version) - - processed = gen.process(wrapper) - - y = yaml.dump(processed) - - output(fname, y) - - # Placeholder for the private.json file. Won't put actual credentials - # here. - output("docker-compose/vertexai/private.json", private_json) - - # Grafana config - with open("grafana/dashboards/dashboard.json") as f: - output( - "docker-compose/grafana/dashboards/dashboard.json", f.read() - ) - - with open("grafana/provisioning/dashboard.yml") as f: - output( - "docker-compose/grafana/provisioning/dashboard.yml", f.read() - ) - - with open("grafana/provisioning/datasource.yml") as f: - output( - "docker-compose/grafana/provisioning/datasource.yml", f.read() - ) - - # Prometheus config - with open("prometheus/prometheus.yml") as f: - output("docker-compose/prometheus/prometheus.yml", f.read()) + return main() diff --git a/templates/generate-all b/templates/generate-all index 1ae71402..22c9a5b0 100755 --- a/templates/generate-all +++ b/templates/generate-all @@ -87,8 +87,18 @@ def full_config_object( ): return config_object([ - graph_store, "pulsar", vector_store, embeddings, - "graph-rag", "grafana", "trustgraph", llm, "workbench-ui", + "triple-store-" + graph_store, + "pulsar", + "vector-store-" + vector_store, + embeddings, + "graph-rag", + "grafana", + "trustgraph-base", + llm, + llm + "-rag", + "workbench-ui", + "prompt-template", + "agent-manager-react", ]) def generate_config( diff --git a/templates/stores/cassandra.jsonnet b/templates/stores/cassandra.jsonnet index 0c90421e..2a9d6d7a 100644 --- a/templates/stores/cassandra.jsonnet +++ b/templates/stores/cassandra.jsonnet @@ -13,7 +13,7 @@ local images = import "values/images.jsonnet"; engine.container("cassandra") .with_image(images.cassandra) .with_environment({ - JVM_OPTS: "-Xms300M -Xmx300M", + JVM_OPTS: "-Xms300M -Xmx300M -Dcassandra.skip_wait_for_gossip_to_settle=0", }) .with_limits("1.0", "1000M") .with_reservations("0.5", "1000M") diff --git a/templates/stores/memgraph.jsonnet b/templates/stores/memgraph.jsonnet index 75faf5f0..70ad127a 100644 --- a/templates/stores/memgraph.jsonnet +++ b/templates/stores/memgraph.jsonnet @@ -7,6 +7,8 @@ local images = import "values/images.jsonnet"; create:: function(engine) + local vol = engine.volume("memgraph").with_size("20G"); + local container = engine.container("memgraph") .with_image(images.memgraph_mage) @@ -16,7 +18,8 @@ local images = import "values/images.jsonnet"; .with_limits("1.0", "1000M") .with_reservations("0.5", "1000M") .with_port(7474, 7474, "api") - .with_port(7687, 7687, "api2"); + .with_port(7687, 7687, "api2") + .with_volume_mount(vol, "/var/lib/memgraph"); local containerSet = engine.containers( "memgraph", [ container ] @@ -28,6 +31,7 @@ local images = import "values/images.jsonnet"; .with_port(7687, 7687, "api2"); engine.resources([ + vol, containerSet, service, ]) @@ -65,4 +69,3 @@ local images = import "values/images.jsonnet"; }, } - diff --git a/templates/stores/neo4j.jsonnet b/templates/stores/neo4j.jsonnet index 55cccc5f..3a8bb783 100644 --- a/templates/stores/neo4j.jsonnet +++ b/templates/stores/neo4j.jsonnet @@ -14,12 +14,14 @@ local images = import "values/images.jsonnet"; .with_image(images.neo4j) .with_environment({ NEO4J_AUTH: "neo4j/password", + NEO4J_server_memory_pagecache_size: "512m", + NEO4J_server_memory_heap_max__size: "512m", // NEO4J_server_bolt_listen__address: "0.0.0.0:7687", // NEO4J_server_default__listen__address: "0.0.0.0", // NEO4J_server_http_listen__address: "0.0.0.0:7474", }) - .with_limits("1.0", "768M") - .with_reservations("0.5", "768M") + .with_limits("1.0", "1536M") + .with_reservations("0.5", "1536M") .with_port(7474, 7474, "api") .with_port(7687, 7687, "api2") .with_volume_mount(vol, "/data"); diff --git a/templates/values/images.jsonnet b/templates/values/images.jsonnet index 40954289..d515a450 100644 --- a/templates/values/images.jsonnet +++ b/templates/values/images.jsonnet @@ -1,7 +1,7 @@ local version = import "version.jsonnet"; { cassandra: "docker.io/cassandra:4.1.6", - neo4j: "docker.io/neo4j:5.22.0-community-bullseye", + neo4j: "docker.io/neo4j:5.26.0-community-bullseye", pulsar: "docker.io/apachepulsar/pulsar:3.3.1", pulsar_manager: "docker.io/apachepulsar/pulsar-manager:v0.4.0", etcd: "quay.io/coreos/etcd:v3.5.15", @@ -9,10 +9,14 @@ local version = import "version.jsonnet"; milvus: "docker.io/milvusdb/milvus:v2.4.9", prometheus: "docker.io/prom/prometheus:v2.53.2", grafana: "docker.io/grafana/grafana:11.1.4", - trustgraph: "docker.io/trustgraph/trustgraph-flow:" + version, + trustgraph_base: "docker.io/trustgraph/trustgraph-base:" + version, + trustgraph_flow: "docker.io/trustgraph/trustgraph-flow:" + version, + trustgraph_bedrock: "docker.io/trustgraph/trustgraph-bedrock:" + version, + trustgraph_vertexai: "docker.io/trustgraph/trustgraph-vertexai:" + version, + trustgraph_hf: "docker.io/trustgraph/trustgraph-hf:" + version, qdrant: "docker.io/qdrant/qdrant:v1.11.1", memgraph_mage: "docker.io/memgraph/memgraph-mage:1.22-memgraph-2.22", memgraph_lab: "docker.io/memgraph/lab:2.19.1", falkordb: "docker.io/falkordb/falkordb:latest", - "workbench-ui": "docker.io/trustgraph/workbench-ui:0.1.6", + "workbench-ui": "docker.io/trustgraph/workbench-ui:0.2.4", } diff --git a/trustgraph-base/trustgraph/api/api.py b/trustgraph-base/trustgraph/api/api.py index de96499c..24207f32 100644 --- a/trustgraph-base/trustgraph/api/api.py +++ b/trustgraph-base/trustgraph/api/api.py @@ -131,6 +131,35 @@ class Api: except: raise ProtocolException(f"Response not formatted correctly") + def document_rag(self, question): + + # The input consists of a question + input = { + "query": question + } + + url = f"{self.url}document-rag" + + # Invoke the API, input is passed as JSON + resp = requests.post(url, json=input) + + # Should be a 200 status code + if resp.status_code != 200: + raise ProtocolException(f"Status code {resp.status_code}") + + try: + # Parse the response as JSON + object = resp.json() + except: + raise ProtocolException(f"Expected JSON response") + + self.check_error(resp) + + try: + return object["response"] + except: + raise ProtocolException(f"Response not formatted correctly") + def embeddings(self, text): # The input consists of a text block diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index eeaf83a1..521dd3c1 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -1,11 +1,15 @@ from pulsar.schema import JsonSchema +import pulsar from prometheus_client import Histogram, Info, Counter, Enum import time from . base_processor import BaseProcessor from .. exceptions import TooManyRequests +default_rate_limit_retry = 10 +default_rate_limit_timeout = 7200 + class Consumer(BaseProcessor): def __init__(self, **params): @@ -21,11 +25,18 @@ class Consumer(BaseProcessor): super(Consumer, self).__init__(**params) - input_queue = params.get("input_queue") - subscriber = params.get("subscriber") - input_schema = params.get("input_schema") + self.input_queue = params.get("input_queue") + self.subscriber = params.get("subscriber") + self.input_schema = params.get("input_schema") - if input_schema == None: + self.rate_limit_retry = params.get( + "rate_limit_retry", default_rate_limit_retry + ) + self.rate_limit_timeout = params.get( + "rate_limit_timeout", default_rate_limit_timeout + ) + + if self.input_schema == None: raise RuntimeError("input_schema must be specified") if not hasattr(__class__, "request_metric"): @@ -43,17 +54,27 @@ class Consumer(BaseProcessor): 'processing_count', 'Processing count', ["status"] ) + if not hasattr(__class__, "rate_limit_metric"): + __class__.rate_limit_metric = Counter( + 'rate_limit_count', 'Rate limit event count', + ) + __class__.pubsub_metric.info({ - "input_queue": input_queue, - "subscriber": subscriber, - "input_schema": input_schema.__name__, + "input_queue": self.input_queue, + "subscriber": self.subscriber, + "input_schema": self.input_schema.__name__, + "rate_limit_retry": str(self.rate_limit_retry), + "rate_limit_timeout": str(self.rate_limit_timeout), }) self.consumer = self.client.subscribe( - input_queue, subscriber, - schema=JsonSchema(input_schema), + self.input_queue, self.subscriber, + consumer_type=pulsar.ConsumerType.Shared, + schema=JsonSchema(self.input_schema), ) + print("Initialised consumer.", flush=True) + def run(self): __class__.state_metric.state('running') @@ -62,31 +83,61 @@ class Consumer(BaseProcessor): msg = self.consumer.receive() - try: + expiry = time.time() + self.rate_limit_timeout - with __class__.request_metric.time(): - self.handle(msg) + # This loop is for retry on rate-limit / resource limits + while True: - # Acknowledge successful processing of the message - self.consumer.acknowledge(msg) + if time.time() > expiry: - __class__.processing_metric.labels(status="success").inc() + print("Gave up waiting for rate-limit retry", flush=True) - except TooManyRequests: - self.consumer.negative_acknowledge(msg) - print("TooManyRequests: will retry") - __class__.processing_metric.labels(status="rate-limit").inc() - time.sleep(5) - continue + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) + + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break + + try: + + with __class__.request_metric.time(): + self.handle(msg) + + # Acknowledge successful processing of the message + self.consumer.acknowledge(msg) + + __class__.processing_metric.labels(status="success").inc() + + # Break out of retry loop + break + + except TooManyRequests: + + print("TooManyRequests: will retry...", flush=True) + + __class__.rate_limit_metric.inc() + + # Sleep + time.sleep(self.rate_limit_retry) + + # Contine from retry loop, just causes a reprocessing + continue - except Exception as e: + except Exception as e: - print("Exception:", e, flush=True) + print("Exception:", e, flush=True) - # Message failed to be processed - self.consumer.negative_acknowledge(msg) + # Message failed to be processed, this causes it to + # be retried + self.consumer.negative_acknowledge(msg) - __class__.processing_metric.labels(status="error").inc() + __class__.processing_metric.labels(status="error").inc() + + # Break out of retry loop, processes next message + break @staticmethod def add_args(parser, default_input_queue, default_subscriber): @@ -105,3 +156,17 @@ class Consumer(BaseProcessor): help=f'Queue subscriber name (default: {default_subscriber})' ) + parser.add_argument( + '--rate-limit-retry', + type=int, + default=default_rate_limit_retry, + help=f'Rate limit retry (default: {default_rate_limit_retry})' + ) + + parser.add_argument( + '--rate-limit-timeout', + type=int, + default=default_rate_limit_timeout, + help=f'Rate limit timeout (default: {default_rate_limit_timeout})' + ) + diff --git a/trustgraph-base/trustgraph/base/consumer_producer.py b/trustgraph-base/trustgraph/base/consumer_producer.py index cabb7525..be9915ce 100644 --- a/trustgraph-base/trustgraph/base/consumer_producer.py +++ b/trustgraph-base/trustgraph/base/consumer_producer.py @@ -1,111 +1,46 @@ from pulsar.schema import JsonSchema +import pulsar from prometheus_client import Histogram, Info, Counter, Enum import time -from . base_processor import BaseProcessor +from . consumer import Consumer from .. exceptions import TooManyRequests -# FIXME: Derive from consumer? And producer? - -class ConsumerProducer(BaseProcessor): +class ConsumerProducer(Consumer): def __init__(self, **params): - if not hasattr(__class__, "state_metric"): - __class__.state_metric = Enum( - 'processor_state', 'Processor state', - states=['starting', 'running', 'stopped'] - ) - __class__.state_metric.state('starting') + super(ConsumerProducer, self).__init__(**params) - __class__.state_metric.state('starting') - - input_queue = params.get("input_queue") - output_queue = params.get("output_queue") - subscriber = params.get("subscriber") - input_schema = params.get("input_schema") - output_schema = params.get("output_schema") - - if not hasattr(__class__, "request_metric"): - __class__.request_metric = Histogram( - 'request_latency', 'Request latency (seconds)' - ) + self.output_queue = params.get("output_queue") + self.output_schema = params.get("output_schema") if not hasattr(__class__, "output_metric"): __class__.output_metric = Counter( 'output_count', 'Output items created' ) - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration' - ) - - if not hasattr(__class__, "processing_metric"): - __class__.processing_metric = Counter( - 'processing_count', 'Processing count', ["status"] - ) - __class__.pubsub_metric.info({ - "input_queue": input_queue, - "output_queue": output_queue, - "subscriber": subscriber, - "input_schema": input_schema.__name__, - "output_schema": output_schema.__name__, + "input_queue": self.input_queue, + "output_queue": self.output_queue, + "subscriber": self.subscriber, + "input_schema": self.input_schema.__name__, + "output_schema": self.output_schema.__name__, + "rate_limit_retry": str(self.rate_limit_retry), + "rate_limit_timeout": str(self.rate_limit_timeout), }) - super(ConsumerProducer, self).__init__(**params) - - if input_schema == None: - raise RuntimeError("input_schema must be specified") - - if output_schema == None: + if self.output_schema == None: raise RuntimeError("output_schema must be specified") self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(output_schema), + topic=self.output_queue, + schema=JsonSchema(self.output_schema), + chunking_enabled=True, ) - self.consumer = self.client.subscribe( - input_queue, subscriber, - schema=JsonSchema(input_schema), - ) - - def run(self): - - __class__.state_metric.state('running') - - while True: - - msg = self.consumer.receive() - - try: - - with __class__.request_metric.time(): - resp = self.handle(msg) - - # Acknowledge successful processing of the message - self.consumer.acknowledge(msg) - - __class__.processing_metric.labels(status="success").inc() - - except TooManyRequests: - self.consumer.negative_acknowledge(msg) - print("TooManyRequests: will retry") - __class__.processing_metric.labels(status="rate-limit").inc() - time.sleep(5) - continue - - except Exception as e: - - print("Exception:", e, flush=True) - - # Message failed to be processed - self.consumer.negative_acknowledge(msg) - - __class__.processing_metric.labels(status="error").inc() + print("Initialised consumer/producer.") def send(self, msg, properties={}): self.producer.send(msg, properties) @@ -117,19 +52,7 @@ class ConsumerProducer(BaseProcessor): default_output_queue, ): - BaseProcessor.add_args(parser) - - parser.add_argument( - '-i', '--input-queue', - default=default_input_queue, - help=f'Input queue (default: {default_input_queue})' - ) - - parser.add_argument( - '-s', '--subscriber', - default=default_subscriber, - help=f'Queue subscriber name (default: {default_subscriber})' - ) + Consumer.add_args(parser, default_input_queue, default_subscriber) parser.add_argument( '-o', '--output-queue', diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 27d693ee..84d7fc99 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/trustgraph-base/trustgraph/base/producer.py @@ -34,6 +34,7 @@ class Producer(BaseProcessor): self.producer = self.client.create_producer( topic=output_queue, schema=JsonSchema(output_schema), + chunking_enabled=True, ) def send(self, msg, properties={}): diff --git a/trustgraph-base/trustgraph/clients/document_embeddings_client.py b/trustgraph-base/trustgraph/clients/document_embeddings_client.py index c27daab0..14547595 100644 --- a/trustgraph-base/trustgraph/clients/document_embeddings_client.py +++ b/trustgraph-base/trustgraph/clients/document_embeddings_client.py @@ -40,8 +40,12 @@ class DocumentEmbeddingsClient(BaseClient): output_schema=DocumentEmbeddingsResponse, ) - def request(self, vectors, limit=10, timeout=300): + def request( + self, vectors, user="trustgraph", collection="default", + limit=10, timeout=300 + ): return self.call( + user=user, collection=collection, vectors=vectors, limit=limit, timeout=timeout ).documents diff --git a/trustgraph-base/trustgraph/exceptions.py b/trustgraph-base/trustgraph/exceptions.py index 16f9956c..afe72ccc 100644 --- a/trustgraph-base/trustgraph/exceptions.py +++ b/trustgraph-base/trustgraph/exceptions.py @@ -8,7 +8,3 @@ class LlmError(Exception): class ParseError(Exception): pass - - - - diff --git a/trustgraph-base/trustgraph/schema/documents.py b/trustgraph-base/trustgraph/schema/documents.py index 2a3d3d0c..fd0049ee 100644 --- a/trustgraph-base/trustgraph/schema/documents.py +++ b/trustgraph-base/trustgraph/schema/documents.py @@ -35,14 +35,18 @@ chunk_ingest_queue = topic('chunk-load') ############################################################################ -# Chunk embeddings are an embeddings associated with a text chunk +# Document embeddings are embeddings associated with a chunk class ChunkEmbeddings(Record): - metadata = Metadata() - vectors = Array(Array(Double())) chunk = Bytes() + vectors = Array(Array(Double())) -chunk_embeddings_ingest_queue = topic('chunk-embeddings-load') +# This is a 'batching' mechanism for the above data +class DocumentEmbeddings(Record): + metadata = Metadata() + chunks = Array(ChunkEmbeddings()) + +document_embeddings_store_queue = topic('document-embeddings-store') ############################################################################ @@ -51,6 +55,8 @@ chunk_embeddings_ingest_queue = topic('chunk-embeddings-load') class DocumentEmbeddingsRequest(Record): vectors = Array(Array(Double())) limit = Integer() + user = String() + collection = String() class DocumentEmbeddingsResponse(Record): error = Error() @@ -62,3 +68,4 @@ document_embeddings_request_queue = topic( document_embeddings_response_queue = topic( 'doc-embeddings', kind='non-persistent', namespace='response', ) + diff --git a/trustgraph-base/trustgraph/schema/graph.py b/trustgraph-base/trustgraph/schema/graph.py index 78c1a99c..7c304e1d 100644 --- a/trustgraph-base/trustgraph/schema/graph.py +++ b/trustgraph-base/trustgraph/schema/graph.py @@ -7,12 +7,31 @@ from . metadata import Metadata ############################################################################ +# Entity context are an entity associated with textual context + +class EntityContext(Record): + entity = Value() + context = String() + +# This is a 'batching' mechanism for the above data +class EntityContexts(Record): + metadata = Metadata() + entities = Array(EntityContext()) + +entity_contexts_ingest_queue = topic('entity-contexts-load') + +############################################################################ + # Graph embeddings are embeddings associated with a graph entity +class EntityEmbeddings(Record): + entity = Value() + vectors = Array(Array(Double())) + +# This is a 'batching' mechanism for the above data class GraphEmbeddings(Record): metadata = Metadata() - vectors = Array(Array(Double())) - entity = Value() + entities = Array(EntityEmbeddings()) graph_embeddings_store_queue = topic('graph-embeddings-store') diff --git a/trustgraph-bedrock/setup.py b/trustgraph-bedrock/setup.py index b8dd36bd..8db4520b 100644 --- a/trustgraph-bedrock/setup.py +++ b/trustgraph-bedrock/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.21,<0.22", "pulsar-client", "prometheus-client", "boto3", diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index a9c05cc8..f46e35fb 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -24,32 +24,48 @@ default_subscriber = module default_model = 'mistral.mistral-large-2407-v1:0' default_temperature = 0.0 default_max_output = 2048 -default_aws_id_key = os.getenv("AWS_ID_KEY", None) -default_aws_secret = os.getenv("AWS_SECRET", None) -default_aws_region = os.getenv("AWS_REGION", 'us-west-2') + +# Actually, these could all just be None, no need to get environment +# variables, as Boto3 would pick all these up if not passed in as args +default_access_key_id = os.getenv("AWS_ACCESS_KEY_ID", None) +default_secret_access_key = os.getenv("AWS_SECRET_ACCESS_KEY", None) +default_session_token = os.getenv("AWS_SESSION_TOKEN", None) +default_profile = os.getenv("AWS_PROFILE", None) +default_region = os.getenv("AWS_DEFAULT_REGION", None) class Processor(ConsumerProducer): def __init__(self, **params): + + print(params) input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) + model = params.get("model", default_model) - aws_id_key = params.get("aws_id_key", default_aws_id_key) - aws_secret = params.get("aws_secret", default_aws_secret) - aws_region = params.get("aws_region", default_aws_region) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) - if aws_id_key is None: - raise RuntimeError("AWS ID not specified") + aws_access_key_id = params.get( + "aws_access_key_id", default_access_key_id + ) - if aws_secret is None: - raise RuntimeError("AWS secret not specified") + aws_secret_access_key = params.get( + "aws_secret_access_key", default_secret_access_key + ) - if aws_region is None: - raise RuntimeError("AWS region not specified") + aws_session_token = params.get( + "aws_session_token", default_session_token + ) + + aws_region = params.get( + "aws_region", default_region + ) + + aws_profile = params.get( + "aws_profile", default_profile + ) super(Processor, self).__init__( **params | { @@ -82,9 +98,11 @@ class Processor(ConsumerProducer): self.max_output = max_output self.session = boto3.Session( - aws_access_key_id=aws_id_key, - aws_secret_access_key=aws_secret, - region_name=aws_region + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + profile_name=aws_profile, + region_name=aws_region, ) self.bedrock = self.session.client(service_name='bedrock-runtime') @@ -179,9 +197,6 @@ class Processor(ConsumerProducer): accept = 'application/json' contentType = 'application/json' - # FIXME: Consider catching request limits and raise TooManyRequests - # See https://boto3.amazonaws.com/v1/documentation/api/latest/guide/retries.html - with __class__.text_completion_metric.time(): response = self.bedrock.invoke_model( body=promptbody, modelId=self.model, accept=accept, @@ -243,30 +258,18 @@ class Processor(ConsumerProducer): print("Done.", flush=True) + except self.bedrock.exceptions.ThrottlingException as e: - # FIXME: Wrong exception, don't know what Bedrock throws - # for a rate limit - except TooManyRequests: + print("Hit rate limit:", e, flush=True) - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + + print(type(e)) print(f"Exception: {e}") print("Send error response...", flush=True) @@ -299,21 +302,27 @@ class Processor(ConsumerProducer): ) parser.add_argument( - '-z', '--aws-id-key', - default=default_aws_id_key, - help=f'AWS ID Key' + '-z', '--aws-access-key-id', + default=default_access_key_id, + help=f'AWS access key ID' ) parser.add_argument( - '-k', '--aws-secret', - default=default_aws_secret, - help=f'AWS Secret Key' + '-k', '--aws-secret-access-key', + default=default_secret_access_key, + help=f'AWS secret access key' ) parser.add_argument( '-r', '--aws-region', - default=default_aws_region, - help=f'AWS Region' + default=default_region, + help=f'AWS region' + ) + + parser.add_argument( + '--aws-profile', '--profile', + default=default_profile, + help=f'AWS profile name' ) parser.add_argument( diff --git a/trustgraph-cli/scripts/tg-dump-msgpack b/trustgraph-cli/scripts/tg-dump-msgpack index 18819649..f3b24d73 100755 --- a/trustgraph-cli/scripts/tg-dump-msgpack +++ b/trustgraph-cli/scripts/tg-dump-msgpack @@ -9,6 +9,7 @@ diagnostic utility. import msgpack import sys import argparse +import json def dump(input_file, action): @@ -17,7 +18,7 @@ def dump(input_file, action): unpacker = msgpack.Unpacker(f, raw=False) for unpacked in unpacker: - print(unpacked) + print(json.dumps(unpacked)) def summary(input_file, action): diff --git a/trustgraph-cli/scripts/tg-graph-show b/trustgraph-cli/scripts/tg-graph-show index 488b7dca..fcba6c23 100755 --- a/trustgraph-cli/scripts/tg-graph-show +++ b/trustgraph-cli/scripts/tg-graph-show @@ -6,24 +6,23 @@ Connects to the graph query service and dumps all graph edges. import argparse import os -from trustgraph.clients.triples_query_client import TriplesQueryClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' -def show_graph(pulsar, user, collection, pulsar_api_key=None): +def show_graph(url, user, collection): - tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key) + api = Api(url) - rows = tq.request( - user=user, collection=collection, - s=None, p=None, o=None, limit=10_000_000 + rows = api.triples_query( +# user=user, collection=collection, + s=None, p=None, o=None, limit=10_000, ) for row in rows: - print(row.s.value, row.p.value, row.o.value) + print(row.s, row.p, row.o) def main(): @@ -33,35 +32,30 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) args = parser.parse_args() try: show_graph( - pulsar=args.pulsar_host, user=args.user, + url=args.api_url, + user=args.user, collection=args.collection, pulsar_api_key=args.pulsar_api_key, ) diff --git a/trustgraph-cli/scripts/tg-graph-to-turtle b/trustgraph-cli/scripts/tg-graph-to-turtle index 65c396f2..900af370 100755 --- a/trustgraph-cli/scripts/tg-graph-to-turtle +++ b/trustgraph-cli/scripts/tg-graph-to-turtle @@ -5,38 +5,45 @@ Connects to the graph query service and dumps all graph edges in Turtle format. """ -import argparse -import os -from trustgraph.clients.triples_query_client import TriplesQueryClient import rdflib import io import sys +import argparse +import os -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +from trustgraph.api import Api, Uri -def show_graph(pulsar, pulsar_api_key=None): +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_user = 'trustgraph' +default_collection = 'default' - tq = TriplesQueryClient(pulsar_host=pulsar, pulsar_api_key=pulsar_api_key) +def show_graph(url, user, collection): - rows = tq.request(None, None, None, limit=10_000_000) + api = Api(url) + + rows = api.triples_query( + s=None, p=None, o=None, + limit=10_000) +# user=user, collection=collection, g = rdflib.Graph() for row in rows: - sv = rdflib.term.URIRef(row.s.value) - pv = rdflib.term.URIRef(row.p.value) + sv = rdflib.term.URIRef(row.s) + pv = rdflib.term.URIRef(row.p) - if row.o.is_uri: + if isinstance(row.o, Uri): # Skip malformed URLs with spaces in - if " " in row.o.value: + if " " in row.o: continue - ov = rdflib.term.URIRef(row.o.value) + ov = rdflib.term.URIRef(row.o) + else: - ov = rdflib.term.Literal(row.o.value) + + ov = rdflib.term.Literal(row.o) g.add((sv, pv, ov)) @@ -57,22 +64,32 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--api-url', + default=default_url, + help=f'API URL (default: {default_url})', ) - + parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', + '-U', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-C', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' ) args = parser.parse_args() try: - show_graph(args.pulsar_host, pulsar_api_key=args.pulsar_api_key) + show_graph( + url=args.api_url, + user=args.user, + collection=args.collection + ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-invoke-agent b/trustgraph-cli/scripts/tg-invoke-agent index b23831cb..bb87c2a6 100755 --- a/trustgraph-cli/scripts/tg-invoke-agent +++ b/trustgraph-cli/scripts/tg-invoke-agent @@ -1,17 +1,18 @@ #!/usr/bin/env python3 """ -Uses the GraphRAG service to answer a query +Uses the GraphRAG service to answer a question """ import argparse import os import textwrap +import uuid +import asyncio +import json +from websockets.asyncio.client import connect -from trustgraph.clients.agent_client import AgentClient - -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -28,15 +29,18 @@ def output(text, prefix="> ", width=78): ) print(out) -def query( - pulsar_host, query, user, collection, - plan=None, state=None, verbose=False, pulsar_api_key=None +async def question( + url, question, user, collection, + plan=None, state=None, verbose=False ): - am = AgentClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key) + if not url.endswith("/"): + url += "/" + + url = url + "api/v1/socket" if verbose: - output(wrap(query), "\U00002753 ") + output(wrap(question), "\U00002753 ") print() def think(x): @@ -49,11 +53,43 @@ def query( output(wrap(x), "\U0001f4a1 ") print() - resp = am.request( - question=query, think=think, observe=observe, - ) + mid = str(uuid.uuid4()) - print(resp) + async with connect(url) as ws: + + req = json.dumps({ + "id": mid, + "service": "agent", + "request": { + "question": question, + } + + }) + + await ws.send(req) + + while True: + + msg = await ws.recv() + + obj = json.loads(msg) + + if obj["id"] != mid: + print("Ignore message") + continue + + if "thought" in obj["response"]: + think(obj["response"]["thought"]) + + if "observation" in obj["response"]: + observe(obj["response"]["observation"]) + + if "answer" in obj["response"]: + print(obj["response"]["answer"]) + + if obj["complete"]: break + + await ws.close() def main(): @@ -63,25 +99,25 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-q', '--query', + '-q', '--question', required=True, - help=f'Query to execute', + help=f'Question to answer', ) parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -101,26 +137,21 @@ def main(): action="store_true", help=f'Output thinking/observations' ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) args = parser.parse_args() try: - query( - pulsar_host=args.pulsar_host, - query=args.query, - user=args.user, - collection=args.collection, - plan=args.plan, - state=args.state, - verbose=args.verbose, - pulsar_api_key=args.pulsar_api_key, + asyncio.run( + question( + url=args.url, + question=args.question, + user=args.user, + collection=args.collection, + plan=args.plan, + state=args.state, + verbose=args.verbose, + ) ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-invoke-document-rag b/trustgraph-cli/scripts/tg-invoke-document-rag new file mode 100755 index 00000000..3e0a9422 --- /dev/null +++ b/trustgraph-cli/scripts/tg-invoke-document-rag @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 + +""" +Uses the GraphRAG service to answer a question +""" + +import argparse +import os +from trustgraph.api import Api + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_user = 'trustgraph' +default_collection = 'default' + +def question(url, question, user, collection): + + rag = Api(url) + +# user=user, collection=collection, + resp = rag.document_rag(question=question) + + print(resp) + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-invoke-document-rag', + description=__doc__, + ) + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-q', '--question', + required=True, + help=f'Question to answer', + ) + + parser.add_argument( + '-U', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-C', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + + args = parser.parse_args() + + try: + + question( + url=args.url, + question=args.question, + user=args.user, + collection=args.collection, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + +main() + diff --git a/trustgraph-cli/scripts/tg-invoke-graph-rag b/trustgraph-cli/scripts/tg-invoke-graph-rag new file mode 100755 index 00000000..44c04041 --- /dev/null +++ b/trustgraph-cli/scripts/tg-invoke-graph-rag @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 + +""" +Uses the GraphRAG service to answer a question +""" + +import argparse +import os +from trustgraph.api import Api + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_user = 'trustgraph' +default_collection = 'default' + +def question(url, question, user, collection): + + rag = Api(url) + +# user=user, collection=collection, + resp = rag.graph_rag(question=question) + + print(resp) + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-invoke-graph-rag', + description=__doc__, + ) + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-q', '--question', + required=True, + help=f'Question to answer', + ) + + parser.add_argument( + '-U', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-C', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + + args = parser.parse_args() + + try: + + question( + url=args.url, + question=args.question, + user=args.user, + collection=args.collection, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + +main() + diff --git a/trustgraph-cli/scripts/tg-invoke-llm b/trustgraph-cli/scripts/tg-invoke-llm index 8e9a4bf2..b9bd73e4 100755 --- a/trustgraph-cli/scripts/tg-invoke-llm +++ b/trustgraph-cli/scripts/tg-invoke-llm @@ -8,17 +8,15 @@ and user prompt. Both arguments are required. import argparse import os import json -from trustgraph.clients.llm_client import LlmClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +def query(url, system, prompt): -def query(pulsar_host, system, prompt, pulsar_api_key=None): + api = Api(url) - cli = LlmClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key) - - resp = cli.request(system=system, prompt=prompt) + resp = api.text_completion(system=system, prompt=prompt) print(resp) @@ -30,9 +28,9 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( @@ -46,23 +44,15 @@ def main(): nargs=1, help='LLM prompt e.g. What is 2 + 2?', ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) - args = parser.parse_args() try: query( - pulsar_host=args.pulsar_host, + url=args.url, system=args.system[0], prompt=args.prompt[0], - pulsar_api_key=args.pulsar_api_key, ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-invoke-prompt b/trustgraph-cli/scripts/tg-invoke-prompt index b165dde6..7d54dab0 100755 --- a/trustgraph-cli/scripts/tg-invoke-prompt +++ b/trustgraph-cli/scripts/tg-invoke-prompt @@ -12,17 +12,15 @@ using key=value arguments on the command line, and these replace import argparse import os import json -from trustgraph.clients.prompt_client import PromptClient +from trustgraph.api import Api -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +def query(url, template_id, variables): -def query(pulsar_host, template_id, variables, pulsar_api_key=None): + api = Api(url) - cli = PromptClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key) - - resp = cli.request(id=template_id, variables=variables) + resp = api.prompt(id=template_id, variables=variables) if isinstance(resp, str): print(resp) @@ -37,9 +35,9 @@ def main(): ) parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( @@ -57,13 +55,6 @@ def main(): specified multiple times''', ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) - args = parser.parse_args() variables = {} @@ -79,10 +70,9 @@ specified multiple times''', try: query( - pulsar_host=args.pulsar_host, + url=args.url, template_id=args.id[0], variables=variables, - pulsar_api_key=args.pulsar_api_key, ) except Exception as e: diff --git a/trustgraph-cli/scripts/tg-load-doc-embeds b/trustgraph-cli/scripts/tg-load-doc-embeds new file mode 100755 index 00000000..d445ec5a --- /dev/null +++ b/trustgraph-cli/scripts/tg-load-doc-embeds @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 + +"""This utility takes a knowledge core and loads it into a running TrustGraph +through the API. The knowledge core should be in msgpack format, which is the +default format produce by tg-save-kg-core. +""" + +import aiohttp +import asyncio +import msgpack +import json +import sys +import argparse +import os +import signal + +class Running: + def __init__(self): self.running = True + def get(self): return self.running + def stop(self): self.running = False + +de_counts = 0 + +async def load_de(running, queue, url): + + global de_counts + + async with aiohttp.ClientSession() as session: + + async with session.ws_connect(f"{url}load/document-embeddings") as ws: + + while running.get(): + + try: + msg = await asyncio.wait_for(queue.get(), 1) + + # End of load + if msg is None: + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + msg = { + "metadata": { + "id": msg["m"]["i"], + "metadata": msg["m"]["m"], + "user": msg["m"]["u"], + "collection": msg["m"]["c"], + }, + "chunks": [ + { + "chunk": chunk["c"], + "vectors": chunk["v"], + } + for chunk in msg["c"] + ], + } + + try: + await ws.send_json(msg) + except Exception as e: + print(e) + + de_counts += 1 + +async def stats(running): + + global de_counts + + while running.get(): + + await asyncio.sleep(2) + + print( + f"Graph embeddings: {de_counts:10d}" + ) + +async def loader(running, de_queue, path, format, user, collection): + + if format == "json": + + raise RuntimeError("Not implemented") + + else: + + with open(path, "rb") as f: + + unpacker = msgpack.Unpacker(f, raw=False) + + while running.get(): + + try: + unpacked = unpacker.unpack() + except: + break + + if user: + unpacked["metadata"]["user"] = user + + if collection: + unpacked["metadata"]["collection"] = collection + + if unpacked[0] == "de": + qtype = de_queue + + while running.get(): + + try: + await asyncio.wait_for(qtype.put(unpacked[1]), 0.5) + + # Successful put message, move on + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + if not running.get(): break + + # Put 'None' on end of queue to finish + while running.get(): + + try: + await asyncio.wait_for(de_queue.put(None), 1) + + # Successful put message, move on + break + + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + +async def run(running, **args): + + # Maxsize on queues reduces back-pressure so tg-load-kg-core doesn't + # grow to eat all memory + de_q = asyncio.Queue(maxsize=10) + + load_task = asyncio.create_task( + loader( + running=running, + de_queue=de_q, + path=args["input_file"], format=args["format"], + user=args["user"], collection=args["collection"], + ) + + ) + + de_task = asyncio.create_task( + load_de( + running=running, + queue=de_q, url=args["url"] + "api/v1/" + ) + ) + + stats_task = asyncio.create_task(stats(running)) + + await de_task + + running.stop() + + await load_task + await stats_task + +async def main(running): + + parser = argparse.ArgumentParser( + prog='tg-load-kg-core', + description=__doc__, + ) + + default_url = os.getenv("TRUSTGRAPH_API", "http://localhost:8088/") + default_user = "trustgraph" + collection = "default" + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'TrustGraph API URL (default: {default_url})', + ) + + parser.add_argument( + '-i', '--input-file', + # Make it mandatory, difficult to over-write an existing file + required=True, + help=f'Output file' + ) + + parser.add_argument( + '--format', + default="msgpack", + choices=["msgpack", "json"], + help=f'Output format (default: msgpack)', + ) + + parser.add_argument( + '--user', + help=f'User ID to load as (default: from input)' + ) + + parser.add_argument( + '--collection', + help=f'Collection ID to load as (default: from input)' + ) + + args = parser.parse_args() + + await run(running, **vars(args)) + +running = Running() + +def interrupt(sig, frame): + running.stop() + print('Interrupt') + +signal.signal(signal.SIGINT, interrupt) + +asyncio.run(main(running)) + diff --git a/trustgraph-cli/scripts/tg-load-kg-core b/trustgraph-cli/scripts/tg-load-kg-core index 4e76e525..b79ec992 100755 --- a/trustgraph-cli/scripts/tg-load-kg-core +++ b/trustgraph-cli/scripts/tg-load-kg-core @@ -51,8 +51,13 @@ async def load_ge(running, queue, url): "user": msg["m"]["u"], "collection": msg["m"]["c"], }, - "vectors": msg["v"], - "entity": msg["e"], + "entities": [ + { + "entity": ent["e"], + "vectors": ent["v"], + } + for ent in msg["e"] + ], } try: diff --git a/trustgraph-cli/scripts/tg-load-pdf b/trustgraph-cli/scripts/tg-load-pdf index 0cdc909f..92140906 100755 --- a/trustgraph-cli/scripts/tg-load-pdf +++ b/trustgraph-cli/scripts/tg-load-pdf @@ -6,21 +6,19 @@ Loads a PDF document into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -import base64 import hashlib import argparse import os import time import uuid -from trustgraph.schema import Document, document_ingest_queue -from trustgraph.schema import Metadata, Triple, Value -from trustgraph.log_level import LogLevel -from trustgraph.knowledge import hash, to_uri, Uri +from trustgraph.api import Api +from trustgraph.knowledge import hash, to_uri from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG from trustgraph.knowledge import Organization, PublicationEvent from trustgraph.knowledge import DigitalDocument +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -28,33 +26,13 @@ class Loader: def __init__( self, - pulsar_host, - output_queue, + url, user, collection, - log_level, metadata, - pulsar_api_key=None, ): - - if pulsar_api_key: - auth = pulsar.AuthenticationToken(pulsar_api_key) - self.client = pulsar.Client( - pulsar_host, - authentication=auth, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - else: - self.client = pulsar.Client( - pulsar_host, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(Document), - chunking_enabled=True, - ) + + self.api = Api(url) self.user = user self.collection = collection @@ -77,49 +55,18 @@ class Loader: id = to_uri(PREF_DOC, id) - triples = [] - - def emit(t): - triples.append(t) - self.metadata.id = id - self.metadata.emit(emit) - r = Document( - metadata=Metadata( - id=id, - metadata=[ - Triple( - s=Value( - value=t["s"], - is_uri=isinstance(t["s"], Uri) - ), - p=Value( - value=t["p"], - is_uri=isinstance(t["p"], Uri) - ), - o=Value( - value=t["o"], - is_uri=isinstance(t["o"], Uri) - ), - ) - for t in triples - ], - user=self.user, - collection=self.collection, - ), - data=base64.b64encode(data), + self.api.load_document( + document=data, id=id, metadata=self.metadata, +# user=self.user, +# collection=self.collection, ) - self.producer.send(r) - print(f"{file}: Loaded successfully.") except Exception as e: print(f"{file}: Failed: {str(e)}", flush=True) - - def __del__(self): - self.client.close() def main(): @@ -128,36 +75,20 @@ def main(): description=__doc__, ) - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') - default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) - default_output_queue = document_ingest_queue - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) - - parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) - - parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -199,7 +130,7 @@ def main(): ) parser.add_argument( - '--url', help=f'Document URL' + '--document-url', help=f'Document URL' ) parser.add_argument( @@ -210,14 +141,6 @@ def main(): '--identifier', '--id', help=f'Document ID' ) - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.ERROR, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - parser.add_argument( 'files', nargs='+', help=f'File to load' @@ -237,7 +160,7 @@ def main(): copyright_holder=args.copyright_holder, copyright_year=args.copyright_year, license=args.license, - url=args.url, + url=args.document_url, keywords=args.keyword, ) @@ -255,12 +178,9 @@ def main(): ) p = Loader( - pulsar_host=args.pulsar_host, - pulsar_api_key=args.pulsar_api_key, - output_queue=args.output_queue, + url=args.url, user=args.user, collection=args.collection, - log_level=args.log_level, metadata=document, ) diff --git a/trustgraph-cli/scripts/tg-load-text b/trustgraph-cli/scripts/tg-load-text index c9dda8a8..a2210dbb 100755 --- a/trustgraph-cli/scripts/tg-load-text +++ b/trustgraph-cli/scripts/tg-load-text @@ -12,14 +12,13 @@ import os import time import uuid -from trustgraph.schema import TextDocument, text_ingest_queue -from trustgraph.schema import Metadata, Triple, Value -from trustgraph.log_level import LogLevel -from trustgraph.knowledge import hash, to_uri, Literal, Uri +from trustgraph.api import Api +from trustgraph.knowledge import hash, to_uri from trustgraph.knowledge import PREF_PUBEV, PREF_DOC, PREF_ORG from trustgraph.knowledge import Organization, PublicationEvent from trustgraph.knowledge import DigitalDocument +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' @@ -27,32 +26,14 @@ class Loader: def __init__( self, - pulsar_host, - output_queue, + url, user, collection, - log_level, metadata, pulsar_api_key=None, ): - if pulsar_api_key: - auth = pulsar.AuthenticationToken(pulsar_api_key) - self.client = pulsar.Client( - pulsar_host, - authentication=auth, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - else: - self.client = pulsar.Client( - pulsar_host, - logger=pulsar.ConsoleLogger(log_level.to_pulsar()) - ) - self.producer = self.client.create_producer( - topic=output_queue, - schema=JsonSchema(TextDocument), - chunking_enabled=True, - ) + self.api = Api(url) self.user = user self.collection = collection @@ -75,49 +56,18 @@ class Loader: id = to_uri(PREF_DOC, id) - triples = [] - - def emit(t): - triples.append(t) - self.metadata.id = id - self.metadata.emit(emit) - r = TextDocument( - metadata=Metadata( - id=id, - metadata=[ - Triple( - s=Value( - value=t["s"], - is_uri=isinstance(t["s"], Uri) - ), - p=Value( - value=t["p"], - is_uri=isinstance(t["p"], Uri) - ), - o=Value( - value=t["o"], - is_uri=isinstance(t["o"], Uri) - ), - ) - for t in triples - ], - user=self.user, - collection=self.collection, - ), - text=data, + self.api.load_text( + text=data, id=id, metadata=self.metadata, +# user=self.user, +# collection=self.collection, ) - self.producer.send(r) - print(f"{file}: Loaded successfully.") except Exception as e: print(f"{file}: Failed: {str(e)}", flush=True) - - def __del__(self): - self.client.close() def main(): @@ -126,37 +76,20 @@ def main(): description=__doc__, ) - default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') - default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) - - default_output_queue = text_ingest_queue - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', - ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', ) parser.add_argument( - '-o', '--output-queue', - default=default_output_queue, - help=f'Output queue (default: {default_output_queue})' - ) - - parser.add_argument( - '-u', '--user', + '-U', '--user', default=default_user, help=f'User ID (default: {default_user})' ) parser.add_argument( - '-c', '--collection', + '-C', '--collection', default=default_collection, help=f'Collection ID (default: {default_collection})' ) @@ -198,7 +131,7 @@ def main(): ) parser.add_argument( - '--url', help=f'Document URL' + '--document-url', help=f'Document URL' ) parser.add_argument( @@ -209,14 +142,6 @@ def main(): '--identifier', '--id', help=f'Document ID' ) - parser.add_argument( - '-l', '--log-level', - type=LogLevel, - default=LogLevel.ERROR, - choices=list(LogLevel), - help=f'Output queue (default: info)' - ) - parser.add_argument( 'files', nargs='+', help=f'File to load' @@ -236,7 +161,7 @@ def main(): copyright_holder=args.copyright_holder, copyright_year=args.copyright_year, license=args.license, - url=args.url, + url=args.document_url, keywords=args.keyword, ) @@ -254,12 +179,9 @@ def main(): ) p = Loader( - pulsar_host=args.pulsar_host, - pulsar_api_key=args.pulsar_api_key, - output_queue=args.output_queue, + url=args.url, user=args.user, collection=args.collection, - log_level=args.log_level, metadata=document, ) diff --git a/trustgraph-cli/scripts/tg-query-document-rag b/trustgraph-cli/scripts/tg-query-document-rag deleted file mode 100755 index ac4c98ca..00000000 --- a/trustgraph-cli/scripts/tg-query-document-rag +++ /dev/null @@ -1,76 +0,0 @@ -#!/usr/bin/env python3 - -""" -Uses the Document RAG service to answer a query -""" - -import argparse -import os -from trustgraph.clients.document_rag_client import DocumentRagClient - -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) - -default_user = 'trustgraph' -default_collection = 'default' - -def query(pulsar_host, query, user, collection, pulsar_api_key=None): - - rag = DocumentRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key) - resp = rag.request(user=user, collection=collection, query=query) - print(resp) - -def main(): - - parser = argparse.ArgumentParser( - prog='tg-query-document-rag', - description=__doc__, - ) - - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', - ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) - parser.add_argument( - '-q', '--query', - required=True, - help=f'Query to execute', - ) - - parser.add_argument( - '-u', '--user', - default=default_user, - help=f'User ID (default: {default_user})' - ) - - parser.add_argument( - '-c', '--collection', - default=default_collection, - help=f'Collection ID (default: {default_collection})' - ) - - args = parser.parse_args() - - try: - - query( - pulsar_host=args.pulsar_host, - pulsar_api_key=args.pulsar_api_key, - query=args.query, - user=args.user, - collection=args.collection, - ) - - except Exception as e: - - print("Exception:", e, flush=True) - -main() - diff --git a/trustgraph-cli/scripts/tg-query-graph-rag b/trustgraph-cli/scripts/tg-query-graph-rag deleted file mode 100755 index 88ab8142..00000000 --- a/trustgraph-cli/scripts/tg-query-graph-rag +++ /dev/null @@ -1,77 +0,0 @@ -#!/usr/bin/env python3 - -""" -Uses the GraphRAG service to answer a query -""" - -import argparse -import os -from trustgraph.clients.graph_rag_client import GraphRagClient - -default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') -default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) - -default_user = 'trustgraph' -default_collection = 'default' - -def query(pulsar_host, query, user, collection, pulsar_api_key=None): - - rag = GraphRagClient(pulsar_host=pulsar_host, pulsar_api_key=pulsar_api_key) - resp = rag.request(user=user, collection=collection, query=query) - print(resp) - -def main(): - - parser = argparse.ArgumentParser( - prog='tg-graph-query-rag', - description=__doc__, - ) - - parser.add_argument( - '-p', '--pulsar-host', - default=default_pulsar_host, - help=f'Pulsar host (default: {default_pulsar_host})', - ) - - parser.add_argument( - '--pulsar-api-key', - default=default_pulsar_api_key, - help=f'Pulsar API key', - ) - - parser.add_argument( - '-q', '--query', - required=True, - help=f'Query to execute', - ) - - parser.add_argument( - '-u', '--user', - default=default_user, - help=f'User ID (default: {default_user})' - ) - - parser.add_argument( - '-c', '--collection', - default=default_collection, - help=f'Collection ID (default: {default_collection})' - ) - - args = parser.parse_args() - - try: - - query( - pulsar_host=args.pulsar_host, - pulsar_api_key=args.pulsar_api_key, - query=args.query, - user=args.user, - collection=args.collection, - ) - - except Exception as e: - - print("Exception:", e, flush=True) - -main() - diff --git a/trustgraph-cli/scripts/tg-save-doc-embeds b/trustgraph-cli/scripts/tg-save-doc-embeds new file mode 100755 index 00000000..95f8b748 --- /dev/null +++ b/trustgraph-cli/scripts/tg-save-doc-embeds @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 + +""" +This utility connects to a running TrustGraph through the API and creates +a knowledge core from the data streaming through the processing queues. +For completeness of data, tg-save-kg-core should be initiated before data +loading takes place. The default output format, msgpack should be used. +JSON output format is also available - msgpack produces a more compact +representation, which is also more performant to load. +""" + +import aiohttp +import asyncio +import msgpack +import json +import sys +import argparse +import os +import signal + +class Running: + def __init__(self): self.running = True + def get(self): return self.running + def stop(self): self.running = False + +async def fetch_de(running, queue, user, collection, url): + + async with aiohttp.ClientSession() as session: + + de_url = f"{url}stream/document-embeddings" + + async with session.ws_connect(de_url) as ws: + + while running.get(): + + try: + msg = await asyncio.wait_for(ws.receive(), 1) + except: + continue + + if msg.type == aiohttp.WSMsgType.TEXT: + + data = msg.json() + + if user: + if data["metadata"]["user"] != user: + continue + + if collection: + if data["metadata"]["collection"] != collection: + continue + + await queue.put([ + "de", + { + "m": { + "i": data["metadata"]["id"], + "m": data["metadata"]["metadata"], + "u": data["metadata"]["user"], + "c": data["metadata"]["collection"], + }, + "c": [ + { + "c": chunk["chunk"], + "v": chunk["vectors"], + } + for chunk in data["chunks"] + ] + } + ]) + if msg.type == aiohttp.WSMsgType.ERROR: + print("Error") + break + +de_counts = 0 + +async def stats(running): + + global t_counts + global de_counts + + while running.get(): + + await asyncio.sleep(2) + + print( + f"Document embeddings: {de_counts:10d}" + ) + +async def output(running, queue, path, format): + + global t_counts + global de_counts + + with open(path, "wb") as f: + + while running.get(): + + try: + msg = await asyncio.wait_for(queue.get(), 0.5) + except: + # Hopefully it's TimeoutError. Annoying to match since + # it changed in 3.11. + continue + + if format == "msgpack": + f.write(msgpack.packb(msg, use_bin_type=True)) + else: + f.write(json.dumps(msg).encode("utf-8")) + + if msg[0] == "de": + de_counts += 1 + + print("Output file closed") + +async def run(running, **args): + + q = asyncio.Queue() + + de_task = asyncio.create_task( + fetch_de( + running=running, + queue=q, user=args["user"], collection=args["collection"], + url=args["url"] + "api/v1/" + ) + ) + + output_task = asyncio.create_task( + output( + running=running, queue=q, + path=args["output_file"], format=args["format"], + ) + + ) + + stats_task = asyncio.create_task(stats(running)) + + await output_task + await de_task + await stats_task + + print("Exiting") + +async def main(running): + + parser = argparse.ArgumentParser( + prog='tg-save-kg-core', + description=__doc__, + ) + + default_url = os.getenv("TRUSTGRAPH_API", "http://localhost:8088/") + default_user = "trustgraph" + collection = "default" + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'TrustGraph API URL (default: {default_url})', + ) + + parser.add_argument( + '-o', '--output-file', + # Make it mandatory, difficult to over-write an existing file + required=True, + help=f'Output file' + ) + + parser.add_argument( + '--format', + default="msgpack", + choices=["msgpack", "json"], + help=f'Output format (default: msgpack)', + ) + + parser.add_argument( + '--user', + help=f'User ID to filter on (default: no filter)' + ) + + parser.add_argument( + '--collection', + help=f'Collection ID to filter on (default: no filter)' + ) + + args = parser.parse_args() + + await run(running, **vars(args)) + +running = Running() + +def interrupt(sig, frame): + running.stop() + print('Interrupt') + +signal.signal(signal.SIGINT, interrupt) + +asyncio.run(main(running)) + diff --git a/trustgraph-cli/scripts/tg-save-kg-core b/trustgraph-cli/scripts/tg-save-kg-core index e52cd7dc..298f2e84 100755 --- a/trustgraph-cli/scripts/tg-save-kg-core +++ b/trustgraph-cli/scripts/tg-save-kg-core @@ -57,8 +57,13 @@ async def fetch_ge(running, queue, user, collection, url): "u": data["metadata"]["user"], "c": data["metadata"]["collection"], }, - "v": data["vectors"], - "e": data["entity"], + "e": [ + { + "e": ent["entity"], + "v": ent["vectors"], + } + for ent in data["entities"] + ] } ]) if msg.type == aiohttp.WSMsgType.ERROR: diff --git a/trustgraph-cli/setup.py b/trustgraph-cli/setup.py index 8217346f..822ab765 100644 --- a/trustgraph-cli/setup.py +++ b/trustgraph-cli/setup.py @@ -34,29 +34,33 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.21,<0.22", "requests", "pulsar-client", + "aiohttp", "rdflib", "tabulate", "msgpack", + "websockets", ], scripts=[ + "scripts/tg-dump-msgpack", "scripts/tg-graph-show", "scripts/tg-graph-to-turtle", + "scripts/tg-init-pulsar", "scripts/tg-init-pulsar-manager", + "scripts/tg-invoke-agent", + "scripts/tg-invoke-document-rag", + "scripts/tg-invoke-graph-rag", + "scripts/tg-invoke-llm", + "scripts/tg-invoke-prompt", + "scripts/tg-load-kg-core", + "scripts/tg-load-doc-embeds", "scripts/tg-load-pdf", "scripts/tg-load-text", "scripts/tg-load-turtle", - "scripts/tg-query-document-rag", - "scripts/tg-query-graph-rag", - "scripts/tg-init-pulsar", "scripts/tg-processor-state", - "scripts/tg-invoke-agent", - "scripts/tg-invoke-prompt", - "scripts/tg-invoke-llm", "scripts/tg-save-kg-core", - "scripts/tg-load-kg-core", - "scripts/tg-dump-msgpack", + "scripts/tg-save-doc-embeds", ] ) diff --git a/trustgraph-embeddings-hf/setup.py b/trustgraph-embeddings-hf/setup.py index 8febd59b..8cf5beb4 100644 --- a/trustgraph-embeddings-hf/setup.py +++ b/trustgraph-embeddings-hf/setup.py @@ -34,8 +34,8 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", - "trustgraph-flow>=0.18,<0.19", + "trustgraph-base>=0.21,<0.22", + "trustgraph-flow>=0.21,<0.22", "torch", "urllib3", "transformers", diff --git a/trustgraph-flow/scripts/document-embeddings b/trustgraph-flow/scripts/document-embeddings new file mode 100755 index 00000000..26bb85b0 --- /dev/null +++ b/trustgraph-flow/scripts/document-embeddings @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.document_embeddings import run + +run() + diff --git a/trustgraph-flow/scripts/embeddings-fastembed b/trustgraph-flow/scripts/embeddings-fastembed new file mode 100755 index 00000000..e1322269 --- /dev/null +++ b/trustgraph-flow/scripts/embeddings-fastembed @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.fastembed import run + +run() + diff --git a/trustgraph-flow/scripts/embeddings-vectorize b/trustgraph-flow/scripts/embeddings-vectorize deleted file mode 100755 index 3de1e3a9..00000000 --- a/trustgraph-flow/scripts/embeddings-vectorize +++ /dev/null @@ -1,6 +0,0 @@ -#!/usr/bin/env python3 - -from trustgraph.embeddings.vectorize import run - -run() - diff --git a/trustgraph-flow/scripts/graph-embeddings b/trustgraph-flow/scripts/graph-embeddings new file mode 100755 index 00000000..29b1fbf4 --- /dev/null +++ b/trustgraph-flow/scripts/graph-embeddings @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.embeddings.graph_embeddings import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 30ec0170..e56aa694 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -34,58 +34,63 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", - "urllib3", - "rdflib", - "pymilvus", - "langchain", - "langchain-core", - "langchain-text-splitters", - "langchain-community", - "requests", - "cassandra-driver", - "pulsar-client", - "pypdf", - "qdrant-client", - "tabulate", + "trustgraph-base>=0.21,<0.22", + "aiohttp", "anthropic", - "pyyaml", - "prometheus-client", + "cassandra-driver", "cohere", - "openai", - "neo4j", - "tiktoken", + "cryptography", + "falkordb", + "fastembed", "google-generativeai", "ibis", "jsonschema", - "aiohttp", + "langchain", + "langchain-community", + "langchain-core", + "langchain-text-splitters", + "neo4j", + "ollama", + "openai", "pinecone[grpc]", - "falkordb", + "prometheus-client", + "pulsar-client", + "pymilvus", + "pypdf", + "pyyaml", + "qdrant-client", + "rdflib", + "requests", + "tabulate", + "tiktoken", + "urllib3", ], scripts=[ - "scripts/api-gateway", "scripts/agent-manager-react", + "scripts/api-gateway", "scripts/chunker-recursive", "scripts/chunker-token", "scripts/de-query-milvus", - "scripts/de-query-qdrant", "scripts/de-query-pinecone", + "scripts/de-query-qdrant", "scripts/de-write-milvus", - "scripts/de-write-qdrant", "scripts/de-write-pinecone", + "scripts/de-write-qdrant", + "scripts/document-embeddings", "scripts/document-rag", "scripts/embeddings-ollama", - "scripts/embeddings-vectorize", + "scripts/embeddings-fastembed", "scripts/ge-query-milvus", "scripts/ge-query-pinecone", "scripts/ge-query-qdrant", "scripts/ge-write-milvus", "scripts/ge-write-pinecone", "scripts/ge-write-qdrant", + "scripts/graph-embeddings", "scripts/graph-rag", "scripts/kg-extract-definitions", - "scripts/kg-extract-topics", "scripts/kg-extract-relationships", + "scripts/kg-extract-topics", "scripts/metering", "scripts/object-extract-row", "scripts/oe-write-milvus", @@ -103,13 +108,13 @@ setuptools.setup( "scripts/text-completion-ollama", "scripts/text-completion-openai", "scripts/triples-query-cassandra", - "scripts/triples-query-neo4j", - "scripts/triples-query-memgraph", "scripts/triples-query-falkordb", + "scripts/triples-query-memgraph", + "scripts/triples-query-neo4j", "scripts/triples-write-cassandra", - "scripts/triples-write-neo4j", - "scripts/triples-write-memgraph", "scripts/triples-write-falkordb", + "scripts/triples-write-memgraph", + "scripts/triples-write-neo4j", "scripts/wikipedia-lookup", ] ) diff --git a/trustgraph-flow/trustgraph/agent/react/service.py b/trustgraph-flow/trustgraph/agent/react/service.py index a5c24563..ccfa1cf8 100755 --- a/trustgraph-flow/trustgraph/agent/react/service.py +++ b/trustgraph-flow/trustgraph/agent/react/service.py @@ -14,8 +14,6 @@ from ... schema import AgentRequest, AgentResponse, AgentStep from ... schema import agent_request_queue, agent_response_queue from ... schema import prompt_request_queue as pr_request_queue from ... schema import prompt_response_queue as pr_response_queue -from ... schema import text_completion_request_queue as tc_request_queue -from ... schema import text_completion_response_queue as tc_response_queue from ... schema import graph_rag_request_queue as gr_request_queue from ... schema import graph_rag_response_queue as gr_response_queue from ... clients.prompt_client import PromptClient @@ -133,12 +131,6 @@ class Processor(ConsumerProducer): prompt_response_queue = params.get( "prompt_response_queue", pr_response_queue ) - text_completion_request_queue = params.get( - "text_completion_request_queue", tc_request_queue - ) - text_completion_response_queue = params.get( - "text_completion_response_queue", tc_response_queue - ) graph_rag_request_queue = params.get( "graph_rag_request_queue", gr_request_queue ) @@ -155,8 +147,6 @@ class Processor(ConsumerProducer): "output_schema": AgentResponse, "prompt_request_queue": prompt_request_queue, "prompt_response_queue": prompt_response_queue, - "text_completion_request_queue": tc_request_queue, - "text_completion_response_queue": tc_response_queue, "graph_rag_request_queue": gr_request_queue, "graph_rag_response_queue": gr_response_queue, } @@ -170,14 +160,6 @@ class Processor(ConsumerProducer): pulsar_api_key=self.pulsar_api_key, ) - self.llm = LlmClient( - subscriber=subscriber, - input_queue=text_completion_request_queue, - output_queue=text_completion_response_queue, - pulsar_host = self.pulsar_host, - pulsar_api_key=self.pulsar_api_key, - ) - self.graph_rag = GraphRagClient( subscriber=subscriber, input_queue=graph_rag_request_queue, @@ -343,18 +325,6 @@ class Processor(ConsumerProducer): help=f'Prompt response queue (default: {pr_response_queue})', ) - parser.add_argument( - '--text-completion-request-queue', - default=tc_request_queue, - help=f'Text completion request queue (default: {tc_request_queue})', - ) - - parser.add_argument( - '--text-completion-response-queue', - default=tc_response_queue, - help=f'Text completion response queue (default: {tc_response_queue})', - ) - parser.add_argument( '--graph-rag-request-queue', default=gr_request_queue, diff --git a/trustgraph-flow/trustgraph/agent/react/tools.py b/trustgraph-flow/trustgraph/agent/react/tools.py index d9bc846f..941610be 100644 --- a/trustgraph-flow/trustgraph/agent/react/tools.py +++ b/trustgraph-flow/trustgraph/agent/react/tools.py @@ -14,6 +14,6 @@ class TextCompletionImpl: self.context = context def invoke(self, **arguments): return self.context.prompt.request( - "question", { "question": arguments.get("computation") } + "question", { "question": arguments.get("question") } ) diff --git a/trustgraph-flow/trustgraph/direct/cassandra.py b/trustgraph-flow/trustgraph/direct/cassandra.py index 568411a9..d06b270f 100644 --- a/trustgraph-flow/trustgraph/direct/cassandra.py +++ b/trustgraph-flow/trustgraph/direct/cassandra.py @@ -6,7 +6,7 @@ class TrustGraph: def __init__( self, hosts=None, - keyspace="trustgraph", table="default", + keyspace="trustgraph", table="default", username=None, password=None ): if hosts is None: @@ -14,8 +14,13 @@ class TrustGraph: self.keyspace = keyspace self.table = table + self.username = username - self.cluster = Cluster(hosts) + if username and password: + auth_provider = PlainTextAuthProvider(username=username, password=password) + self.cluster = Cluster(hosts, auth_provider=auth_provider) + else: + self.cluster = Cluster(hosts) self.session = self.cluster.connect() self.init() diff --git a/trustgraph-flow/trustgraph/document_rag.py b/trustgraph-flow/trustgraph/document_rag.py index 5c340150..f4676b15 100644 --- a/trustgraph-flow/trustgraph/document_rag.py +++ b/trustgraph-flow/trustgraph/document_rag.py @@ -16,6 +16,44 @@ from . schema import document_embeddings_response_queue LABEL="http://www.w3.org/2000/01/rdf-schema#label" DEFINITION="http://www.w3.org/2004/02/skos/core#definition" +class Query: + + def __init__(self, rag, user, collection, verbose): + self.rag = rag + self.user = user + self.collection = collection + self.verbose = verbose + + def get_vector(self, query): + + if self.verbose: + print("Compute embeddings...", flush=True) + + qembeds = self.rag.embeddings.request(query) + + if self.verbose: + print("Done.", flush=True) + + return qembeds + + def get_docs(self, query): + + vectors = self.get_vector(query) + + if self.verbose: + print("Get entities...", flush=True) + + docs = self.rag.de_client.request( + vectors, limit=self.rag.doc_limit + ) + + if self.verbose: + print("Docs:", flush=True) + for doc in docs: + print(doc, flush=True) + + return docs + class DocumentRag: def __init__( @@ -56,7 +94,7 @@ class DocumentRag: print("Initialising...", flush=True) # FIXME: Configurable - self.entity_limit = 20 + self.doc_limit = 20 self.de_client = DocumentEmbeddingsClient( pulsar_host=pulsar_host, @@ -85,42 +123,16 @@ class DocumentRag: if self.verbose: print("Initialised", flush=True) - def get_vector(self, query): - - if self.verbose: - print("Compute embeddings...", flush=True) - - qembeds = self.embeddings.request(query) - - if self.verbose: - print("Done.", flush=True) - - return qembeds - - def get_docs(self, query): - - vectors = self.get_vector(query) - - if self.verbose: - print("Get entities...", flush=True) - - docs = self.de_client.request( - vectors, self.entity_limit - ) - - if self.verbose: - print("Docs:", flush=True) - for doc in docs: - print(doc, flush=True) - - return docs - - def query(self, query): + def query(self, query, user="trustgraph", collection="default"): if self.verbose: print("Construct prompt...", flush=True) - docs = self.get_docs(query) + q = Query( + rag=self, user=user, collection=collection, verbose=self.verbose + ) + + docs = q.get_docs(query) if self.verbose: print("Invoke LLM...", flush=True) diff --git a/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py new file mode 100644 index 00000000..40d505a5 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__init__.py @@ -0,0 +1,3 @@ + +from . embeddings import * + diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py similarity index 57% rename from trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py rename to trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py index a578de8a..a48cc4d0 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/__main__.py +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/__main__.py @@ -1,5 +1,5 @@ -from . vectorize import run +from . embeddings import run if __name__ == '__main__': run() diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py similarity index 75% rename from trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py rename to trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py index c24b8251..48486378 100755 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py @@ -1,11 +1,13 @@ """ -Vectorizer, calls the embeddings service to get embeddings for a chunk. -Input is text chunk, output is chunk and vectors. +Document embeddings, calls the embeddings service to get embeddings for a +chunk of text. Input is chunk of text plus metadata. +Output is chunk plus embedding. """ -from ... schema import Chunk, ChunkEmbeddings -from ... schema import chunk_ingest_queue, chunk_embeddings_ingest_queue +from ... schema import Chunk, ChunkEmbeddings, DocumentEmbeddings +from ... schema import chunk_ingest_queue +from ... schema import document_embeddings_store_queue from ... schema import embeddings_request_queue, embeddings_response_queue from ... clients.embeddings_client import EmbeddingsClient from ... log_level import LogLevel @@ -14,7 +16,7 @@ from ... base import ConsumerProducer module = ".".join(__name__.split(".")[1:-1]) default_input_queue = chunk_ingest_queue -default_output_queue = chunk_embeddings_ingest_queue +default_output_queue = document_embeddings_store_queue default_subscriber = module class Processor(ConsumerProducer): @@ -39,7 +41,7 @@ class Processor(ConsumerProducer): "embeddings_response_queue": emb_response_queue, "subscriber": subscriber, "input_schema": Chunk, - "output_schema": ChunkEmbeddings, + "output_schema": DocumentEmbeddings, } ) @@ -51,31 +53,35 @@ class Processor(ConsumerProducer): subscriber=module + "-emb", ) - def emit(self, metadata, chunk, vectors): - - r = ChunkEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors) - self.producer.send(r) - def handle(self, msg): v = msg.value() print(f"Indexing {v.metadata.id}...", flush=True) - chunk = v.chunk.decode("utf-8") - try: - vectors = self.embeddings.request(chunk) + vectors = self.embeddings.request(v.chunk) - self.emit( + embeds = [ + ChunkEmbeddings( + chunk=v.chunk, + vectors=vectors, + ) + ] + + r = DocumentEmbeddings( metadata=v.metadata, - chunk=chunk.encode("utf-8"), - vectors=vectors + chunks=embeds, ) + self.producer.send(r) + except Exception as e: print("Exception:", e, flush=True) + # Retry + raise e + print("Done.", flush=True) @staticmethod diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py b/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py new file mode 100644 index 00000000..9d16af90 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/__init__.py @@ -0,0 +1,3 @@ + +from . processor import * + diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py b/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py new file mode 100755 index 00000000..986c0257 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . processor import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py new file mode 100755 index 00000000..635387b8 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -0,0 +1,89 @@ + +""" +Embeddings service, applies an embeddings model selected from HuggingFace. +Input is text, output is embeddings vector. +""" + +from ... schema import EmbeddingsRequest, EmbeddingsResponse +from ... schema import embeddings_request_queue, embeddings_response_queue +from ... log_level import LogLevel +from ... base import ConsumerProducer +from fastembed import TextEmbedding +import os + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = embeddings_request_queue +default_output_queue = embeddings_response_queue +default_subscriber = module +default_model="sentence-transformers/all-MiniLM-L6-v2" + +class Processor(ConsumerProducer): + + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + + model = params.get("model", default_model) + + super(Processor, self).__init__( + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": EmbeddingsRequest, + "output_schema": EmbeddingsResponse, + "model": model, + } + ) + + self.embeddings = TextEmbedding(model_name = model) + + def handle(self, msg): + + v = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + text = v.text + vecs = self.embeddings.embed([text]) + + vecs = [ + v.tolist() + for v in vecs + ] + + print("Send response...", flush=True) + r = EmbeddingsResponse( + vectors=list(vecs), + error=None, + ) + + self.producer.send(r, properties={"id": id}) + + print("Done.", flush=True) + + @staticmethod + def add_args(parser): + + ConsumerProducer.add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ) + + parser.add_argument( + '-m', '--model', + default=default_model, + help=f'Embeddings model (default: {default_model})' + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py new file mode 100644 index 00000000..40d505a5 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__init__.py @@ -0,0 +1,3 @@ + +from . embeddings import * + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py new file mode 100755 index 00000000..a48cc4d0 --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/__main__.py @@ -0,0 +1,6 @@ + +from . embeddings import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py new file mode 100755 index 00000000..e4d1646e --- /dev/null +++ b/trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py @@ -0,0 +1,113 @@ + +""" +Graph embeddings, calls the embeddings service to get embeddings for a +set of entity contexts. Input is entity plus textual context. +Output is entity plus embedding. +""" + +from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings +from ... schema import entity_contexts_ingest_queue +from ... schema import graph_embeddings_store_queue +from ... schema import embeddings_request_queue, embeddings_response_queue +from ... clients.embeddings_client import EmbeddingsClient +from ... log_level import LogLevel +from ... base import ConsumerProducer + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = entity_contexts_ingest_queue +default_output_queue = graph_embeddings_store_queue +default_subscriber = module + +class Processor(ConsumerProducer): + + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + emb_request_queue = params.get( + "embeddings_request_queue", embeddings_request_queue + ) + emb_response_queue = params.get( + "embeddings_response_queue", embeddings_response_queue + ) + + super(Processor, self).__init__( + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "embeddings_request_queue": emb_request_queue, + "embeddings_response_queue": emb_response_queue, + "subscriber": subscriber, + "input_schema": EntityContexts, + "output_schema": GraphEmbeddings, + } + ) + + self.embeddings = EmbeddingsClient( + pulsar_host=self.pulsar_host, + input_queue=emb_request_queue, + output_queue=emb_response_queue, + subscriber=module + "-emb", + ) + + def handle(self, msg): + + v = msg.value() + print(f"Indexing {v.metadata.id}...", flush=True) + + entities = [] + + try: + + for entity in v.entities: + + vectors = self.embeddings.request(entity.context) + + entities.append( + EntityEmbeddings( + entity=entity.entity, + vectors=vectors + ) + ) + + r = GraphEmbeddings( + metadata=v.metadata, + entities=entities, + ) + + self.producer.send(r) + + except Exception as e: + print("Exception:", e, flush=True) + + # Retry + raise e + + print("Done.", flush=True) + + @staticmethod + def add_args(parser): + + ConsumerProducer.add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ) + + parser.add_argument( + '--embeddings-request-queue', + default=embeddings_request_queue, + help=f'Embeddings request queue (default: {embeddings_request_queue})', + ) + + parser.add_argument( + '--embeddings-response-queue', + default=embeddings_response_queue, + help=f'Embeddings request queue (default: {embeddings_response_queue})', + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index 6682a79f..5baf64aa 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -1,14 +1,15 @@ """ -Embeddings service, applies an embeddings model selected from HuggingFace. +Embeddings service, applies an embeddings model hosted on a local Ollama. Input is text, output is embeddings vector. """ -from langchain_community.embeddings import OllamaEmbeddings from ... schema import EmbeddingsRequest, EmbeddingsResponse from ... schema import embeddings_request_queue, embeddings_response_queue from ... log_level import LogLevel from ... base import ConsumerProducer +from ollama import Client +import os module = ".".join(__name__.split(".")[1:-1]) @@ -16,7 +17,7 @@ default_input_queue = embeddings_request_queue default_output_queue = embeddings_response_queue default_subscriber = module default_model="mxbai-embed-large" -default_ollama = 'http://localhost:11434' +default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434') class Processor(ConsumerProducer): @@ -26,6 +27,9 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) + ollama = params.get("ollama", default_ollama) + model = params.get("model", default_model) + super(Processor, self).__init__( **params | { "input_queue": input_queue, @@ -33,10 +37,13 @@ class Processor(ConsumerProducer): "subscriber": subscriber, "input_schema": EmbeddingsRequest, "output_schema": EmbeddingsResponse, + "ollama": ollama, + "model": model, } ) - self.embeddings = OllamaEmbeddings(base_url=ollama, model=model) + self.client = Client(host=ollama) + self.model = model def handle(self, msg): @@ -49,10 +56,16 @@ class Processor(ConsumerProducer): print(f"Handling input {id}...", flush=True) text = v.text - embeds = self.embeddings.embed_query([text]) + embeds = self.client.embed( + model = self.model, + input = text + ) print("Send response...", flush=True) - r = EmbeddingsResponse(vectors=[embeds]) + r = EmbeddingsResponse( + vectors=embeds.embeddings, + error=None, + ) self.producer.send(r, properties={"id": id}) diff --git a/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py b/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py deleted file mode 100644 index 31596b8c..00000000 --- a/trustgraph-flow/trustgraph/embeddings/vectorize/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ - -from . vectorize import * - diff --git a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py index 7b82e0d7..fe824378 100755 --- a/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/definitions/extract.py @@ -1,14 +1,17 @@ """ -Simple decoder, accepts embeddings+text chunks input, applies entity analysis to -get entity definitions which are output as graph edges. +Simple decoder, accepts text chunks input, applies entity analysis to +get entity definitions which are output as graph edges along with +entity/context definitions for embedding. """ import urllib.parse -import json +from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue +from .... schema import Chunk, Triple, Triples, Metadata, Value +from .... schema import EntityContext, EntityContexts +from .... schema import chunk_ingest_queue, triples_store_queue +from .... schema import entity_contexts_ingest_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue from .... log_level import LogLevel @@ -22,8 +25,9 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = chunk_ingest_queue default_output_queue = triples_store_queue +default_entity_context_queue = entity_contexts_ingest_queue default_subscriber = module class Processor(ConsumerProducer): @@ -32,6 +36,10 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) + ec_queue = params.get( + "entity_context_queue", + default_entity_context_queue + ) subscriber = params.get("subscriber", default_subscriber) pr_request_queue = params.get( "prompt_request_queue", prompt_request_queue @@ -45,13 +53,30 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": Chunk, "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } ) + self.ec_prod = self.client.create_producer( + topic=ec_queue, + schema=JsonSchema(EntityContexts), + ) + + __class__.pubsub_metric.info({ + "input_queue": input_queue, + "output_queue": output_queue, + "entity_context_queue": ec_queue, + "prompt_request_queue": pr_request_queue, + "prompt_response_queue": pr_response_queue, + "subscriber": subscriber, + "input_schema": Chunk.__name__, + "output_schema": Triples.__name__, + "vector_schema": EntityContexts.__name__, + }) + self.prompt = PromptClient( pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key, @@ -80,6 +105,14 @@ class Processor(ConsumerProducer): ) self.producer.send(t) + def emit_ecs(self, metadata, entities): + + t = EntityContexts( + metadata=metadata, + entities=entities, + ) + self.ec_prod.send(t) + def handle(self, msg): v = msg.value() @@ -92,6 +125,7 @@ class Processor(ConsumerProducer): defs = self.get_definitions(chunk) triples = [] + entities = [] # FIXME: Putting metadata into triples store is duplicated in # relationships extractor too @@ -130,6 +164,14 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) + ec = EntityContext( + entity=s_value, + context=defn.definition, + ) + + entities.append(ec) + + self.emit_edges( Metadata( id=v.metadata.id, @@ -140,6 +182,16 @@ class Processor(ConsumerProducer): triples ) + self.emit_ecs( + Metadata( + id=v.metadata.id, + metadata=[], + user=v.metadata.user, + collection=v.metadata.collection, + ), + entities + ) + except Exception as e: print("Exception: ", e, flush=True) @@ -153,6 +205,12 @@ class Processor(ConsumerProducer): default_output_queue, ) + parser.add_argument( + '-e', '--entity-context-queue', + default=default_entity_context_queue, + help=f'Entity context queue (default: {default_entity_context_queue})' + ) + parser.add_argument( '--prompt-request-queue', default=prompt_request_queue, diff --git a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py index 34f36ef6..fae6e07b 100755 --- a/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/relationships/extract.py @@ -1,18 +1,15 @@ """ -Simple decoder, accepts vector+text chunks input, applies entity +Simple decoder, accepts text chunks input, applies entity relationship analysis to get entity relationship edges which are output as graph edges. """ import urllib.parse -import os -from pulsar.schema import JsonSchema -from .... schema import ChunkEmbeddings, Triple, Triples, GraphEmbeddings +from .... schema import Chunk, Triple, Triples from .... schema import Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue -from .... schema import graph_embeddings_store_queue +from .... schema import chunk_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue from .... log_level import LogLevel @@ -25,9 +22,8 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = chunk_ingest_queue default_output_queue = triples_store_queue -default_vector_queue = graph_embeddings_store_queue default_subscriber = module class Processor(ConsumerProducer): @@ -36,7 +32,6 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) - vector_queue = params.get("vector_queue", default_vector_queue) subscriber = params.get("subscriber", default_subscriber) pr_request_queue = params.get( "prompt_request_queue", prompt_request_queue @@ -50,30 +45,13 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": Chunk, "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, } ) - self.vec_prod = self.client.create_producer( - topic=vector_queue, - schema=JsonSchema(GraphEmbeddings), - ) - - __class__.pubsub_metric.info({ - "input_queue": input_queue, - "output_queue": output_queue, - "vector_queue": vector_queue, - "prompt_request_queue": pr_request_queue, - "prompt_response_queue": pr_response_queue, - "subscriber": subscriber, - "input_schema": ChunkEmbeddings.__name__, - "output_schema": Triples.__name__, - "vector_schema": GraphEmbeddings.__name__, - }) - self.prompt = PromptClient( pulsar_host=self.pulsar_host, pulsar_api_key=self.pulsar_api_key, @@ -102,11 +80,6 @@ class Processor(ConsumerProducer): ) self.producer.send(t) - def emit_vec(self, metadata, ent, vec): - - r = GraphEmbeddings(metadata=metadata, entity=ent, vectors=vec) - self.vec_prod.send(r) - def handle(self, msg): v = msg.value() @@ -194,12 +167,6 @@ class Processor(ConsumerProducer): o=Value(value=v.metadata.id, is_uri=True) )) - self.emit_vec(v.metadata, s_value, v.vectors) - self.emit_vec(v.metadata, p_value, v.vectors) - - if rel.o_entity: - self.emit_vec(v.metadata, o_value, v.vectors) - self.emit_edges( Metadata( id=v.metadata.id, @@ -223,12 +190,6 @@ class Processor(ConsumerProducer): default_output_queue, ) - parser.add_argument( - '-c', '--vector-queue', - default=default_vector_queue, - help=f'Vector output queue (default: {default_vector_queue})' - ) - parser.add_argument( '--prompt-request-queue', default=prompt_request_queue, diff --git a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py index d1c0c3bc..b68bb3ea 100755 --- a/trustgraph-flow/trustgraph/extract/kg/topics/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/topics/extract.py @@ -1,14 +1,14 @@ """ -Simple decoder, accepts embeddings+text chunks input, applies entity analysis to -get entity definitions which are output as graph edges. +Simple decoder, accepts text chunks input, applies entity analysis to +get topics which are output as graph edges. """ import urllib.parse import json -from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value -from .... schema import chunk_embeddings_ingest_queue, triples_store_queue +from .... schema import Chunk, Triple, Triples, Metadata, Value +from .... schema import chunk_ingest_queue, triples_store_queue from .... schema import prompt_request_queue from .... schema import prompt_response_queue from .... log_level import LogLevel @@ -20,7 +20,7 @@ DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True) module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = chunk_ingest_queue default_output_queue = triples_store_queue default_subscriber = module @@ -43,7 +43,7 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": Chunk, "output_schema": Triples, "prompt_request_queue": pr_request_queue, "prompt_response_queue": pr_response_queue, diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py new file mode 100644 index 00000000..1a7f635d --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py @@ -0,0 +1,64 @@ + +import asyncio +from pulsar.schema import JsonSchema +import uuid +from aiohttp import WSMsgType + +from .. schema import Metadata +from .. schema import DocumentEmbeddings, ChunkEmbeddings +from .. schema import document_embeddings_store_queue + +from . publisher import Publisher +from . socket import SocketEndpoint +from . serialize import to_subgraph + +class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): + + def __init__( + self, pulsar_host, auth, path="/api/v1/load/document-embeddings", + ): + + super(DocumentEmbeddingsLoadEndpoint, self).__init__( + endpoint_path=path, auth=auth, + ) + + self.pulsar_host=pulsar_host + + self.publisher = Publisher( + self.pulsar_host, document_embeddings_store_queue, + schema=JsonSchema(DocumentEmbeddings) + ) + + async def start(self): + + self.publisher.start() + + async def listener(self, ws, running): + + async for msg in ws: + # On error, finish + if msg.type == WSMsgType.ERROR: + break + else: + + data = msg.json() + + elt = DocumentEmbeddings( + metadata=Metadata( + id=data["metadata"]["id"], + metadata=to_subgraph(data["metadata"]["metadata"]), + user=data["metadata"]["user"], + collection=data["metadata"]["collection"], + ), + chunks=[ + ChunkEmbeddings( + chunk=de["chunk"].encode("utf-8"), + vectors=de["vectors"], + ) + for de in data["chunks"] + ], + ) + + self.publisher.send(None, elt) + + running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py new file mode 100644 index 00000000..99cfb0a9 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py @@ -0,0 +1,72 @@ + +import asyncio +import queue +from pulsar.schema import JsonSchema +import uuid + +from .. schema import DocumentEmbeddings +from .. schema import document_embeddings_store_queue + +from . subscriber import Subscriber +from . socket import SocketEndpoint +from . serialize import serialize_document_embeddings + +class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): + + def __init__( + self, pulsar_host, auth, path="/api/v1/stream/document-embeddings" + ): + + super(DocumentEmbeddingsStreamEndpoint, self).__init__( + endpoint_path=path, auth=auth, + ) + + self.pulsar_host=pulsar_host + + self.subscriber = Subscriber( + self.pulsar_host, document_embeddings_store_queue, + "api-gateway", "api-gateway", + schema=JsonSchema(DocumentEmbeddings) + ) + + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(DocumentEmbeddingsStreamEndpoint, self).listener( + ws, running + ) + + await worker + + async def start(self): + + self.subscriber.start() + + async def async_thread(self, ws, running): + + id = str(uuid.uuid4()) + + q = self.subscriber.subscribe_all(id) + + while running.get(): + try: + resp = await asyncio.to_thread(q.get, timeout=0.5) + await ws.send_json(serialize_document_embeddings(resp)) + + except TimeoutError: + continue + + except queue.Empty: + continue + + except Exception as e: + print(f"Exception: {str(e)}", flush=True) + break + + self.subscriber.unsubscribe_all(id) + + running.stop() + diff --git a/trustgraph-flow/trustgraph/gateway/document_load.py b/trustgraph-flow/trustgraph/gateway/document_load.py index 55ee2038..7853c9bd 100644 --- a/trustgraph-flow/trustgraph/gateway/document_load.py +++ b/trustgraph-flow/trustgraph/gateway/document_load.py @@ -1,7 +1,7 @@ import base64 -from .. schema import Document +from .. schema import Document, Metadata from .. schema import document_ingest_queue from . sender import ServiceSender @@ -19,25 +19,24 @@ class DocumentLoadSender(ServiceSender): def to_request(self, body): - if "metadata" in data: - metadata = to_subgraph(data["metadata"]) + if "metadata" in body: + metadata = to_subgraph(body["metadata"]) else: metadata = [] # Doing a base64 decoe/encode here to make sure the # content is valid base64 - doc = base64.b64decode(data["data"]) + doc = base64.b64decode(body["data"]) print("Document received") return Document( metadata=Metadata( - id=data.get("id"), + id=body.get("id"), metadata=metadata, - user=data.get("user", "trustgraph"), - collection=data.get("collection", "default"), + user=body.get("user", "trustgraph"), + collection=body.get("collection", "default"), ), data=base64.b64encode(doc).decode("utf-8") ) - diff --git a/trustgraph-flow/trustgraph/gateway/document_rag.py b/trustgraph-flow/trustgraph/gateway/document_rag.py new file mode 100644 index 00000000..6665a7dc --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/document_rag.py @@ -0,0 +1,30 @@ + +from .. schema import DocumentRagQuery, DocumentRagResponse +from .. schema import document_rag_request_queue +from .. schema import document_rag_response_queue + +from . endpoint import ServiceEndpoint +from . requestor import ServiceRequestor + +class DocumentRagRequestor(ServiceRequestor): + def __init__(self, pulsar_host, timeout, auth): + + super(DocumentRagRequestor, self).__init__( + pulsar_host=pulsar_host, + request_queue=document_rag_request_queue, + response_queue=document_rag_response_queue, + request_schema=DocumentRagQuery, + response_schema=DocumentRagResponse, + timeout=timeout, + ) + + def to_request(self, body): + return DocumentRagQuery( + query=body["query"], + user=body.get("user", "trustgraph"), + collection=body.get("collection", "default"), + ) + + def from_response(self, message): + return { "response": message.response }, True + diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index 236a4ea6..0d4b12e9 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py @@ -5,7 +5,7 @@ import uuid from aiohttp import WSMsgType from .. schema import Metadata -from .. schema import GraphEmbeddings +from .. schema import GraphEmbeddings, EntityEmbeddings from .. schema import graph_embeddings_store_queue from . publisher import Publisher @@ -52,8 +52,13 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): user=data["metadata"]["user"], collection=data["metadata"]["collection"], ), - entity=to_value(data["entity"]), - vectors=data["vectors"], + entities=[ + EntityEmbeddings( + entity=to_value(ent["entity"]), + vectors=ent["vectors"], + ) + for ent in data["entities"] + ] ) self.publisher.send(None, elt) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py index ec7a3f9d..c2613144 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py @@ -31,6 +31,16 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): schema=JsonSchema(GraphEmbeddings) ) + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(GraphEmbeddingsStreamEndpoint, self).listener(ws, running) + + await worker + async def start(self): self.subscriber.start() @@ -46,6 +56,9 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_graph_embeddings(resp)) + except TimeoutError: + continue + except queue.Empty: continue diff --git a/trustgraph-flow/trustgraph/gateway/metrics.py b/trustgraph-flow/trustgraph/gateway/metrics.py new file mode 100644 index 00000000..33c1fe3a --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/metrics.py @@ -0,0 +1,73 @@ + +# +# This provides a Prometheus endpoint on the api-gateway. It proxies +# HTTP GET requests to Prometheus. +# + +import aiohttp +from aiohttp import web +import asyncio +from pulsar.schema import JsonSchema +import uuid +import logging + +logger = logging.getLogger("endpoint") +logger.setLevel(logging.INFO) + +class MetricsEndpoint: + + def __init__(self, prometheus_url, endpoint_path, auth): + + self.prometheus_url = prometheus_url + self.path = endpoint_path + self.auth = auth + self.operation = "service" + + async def start(self): + pass + + def add_routes(self, app): + + app.add_routes([ + web.get(self.path + "/{path:.*}", self.handle), + ]) + + async def handle(self, request): + + print(request.path, "...") + + try: + ht = request.headers["Authorization"] + tokens = ht.split(" ", 2) + if tokens[0] != "Bearer": + return web.HTTPUnauthorized() + token = tokens[1] + except: + token = "" + + if not self.auth.permitted(token, self.operation): + return web.HTTPUnauthorized() + + try: + + path = request.match_info["path"] + + async with aiohttp.ClientSession() as session: + + url = ( + self.prometheus_url + "/api/v1/" + path + "?" + + request.query_string + ) + + async with session.get(url) as resp: + return web.Response( + status=resp.status, + text=await resp.text() + ) + + except Exception as e: + + logging.error(f"Exception: {e}") + + raise web.HTTPInternalServerError() + diff --git a/trustgraph-flow/trustgraph/gateway/publisher.py b/trustgraph-flow/trustgraph/gateway/publisher.py index ab752cba..5dfb882b 100644 --- a/trustgraph-flow/trustgraph/gateway/publisher.py +++ b/trustgraph-flow/trustgraph/gateway/publisher.py @@ -7,13 +7,14 @@ import threading class Publisher: def __init__(self, pulsar_host, topic, schema=None, max_size=10, - chunking_enabled=False, pulsar_api_key=None): + chunking_enabled=True, listener=None, pulsar_api_key=None): self.pulsar_host = pulsar_host self.pulsar_api_key = pulsar_api_key, self.topic = topic self.schema = schema self.q = queue.Queue(maxsize=max_size) self.chunking_enabled = chunking_enabled + self.listener_name = listener def start(self): self.task = threading.Thread(target=self.run) @@ -28,11 +29,13 @@ class Publisher: if self.pulsar_api_key: client = pulsar.Client( self.pulsar_host, + listener_name=self.listener_name, authentication=pulsar.AuthenticationToken(self.pulsar_api_key) ) else: client = pulsar.Client( self.pulsar_host, + listener_name=self.listener_name ) producer = client.create_producer( diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index 8e2911d1..9216f681 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -63,12 +63,18 @@ class ServiceRequestor: while True: try: - resp = await asyncio.to_thread(q.get, timeout=self.timeout) + resp = await asyncio.to_thread( + q.get, + timeout=self.timeout + ) except Exception as e: raise RuntimeError("Timeout") if resp.error: - return { "error": resp.error.message } + err = { "error": resp.error.message } + if responder: + await responder(err, True) + return err resp, fin = self.from_response(resp) @@ -84,7 +90,10 @@ class ServiceRequestor: logging.error(f"Exception: {e}") - return { "error": str(e) } + err = { "error": str(e) } + if responder: + await responder(err, True) + return err finally: self.sub.unsubscribe(id) diff --git a/trustgraph-flow/trustgraph/gateway/sender.py b/trustgraph-flow/trustgraph/gateway/sender.py index 91e3469c..f232780f 100644 --- a/trustgraph-flow/trustgraph/gateway/sender.py +++ b/trustgraph-flow/trustgraph/gateway/sender.py @@ -48,5 +48,11 @@ class ServiceSender: logging.error(f"Exception: {e}") - return { "error": str(e) } + err = { "error": str(e) } + + if responder: + await responder(err, True) + + return err + diff --git a/trustgraph-flow/trustgraph/gateway/serialize.py b/trustgraph-flow/trustgraph/gateway/serialize.py index 35932382..40b6efc5 100644 --- a/trustgraph-flow/trustgraph/gateway/serialize.py +++ b/trustgraph-flow/trustgraph/gateway/serialize.py @@ -51,7 +51,29 @@ def serialize_graph_embeddings(message): "user": message.metadata.user, "collection": message.metadata.collection, }, - "vectors": message.vectors, - "entity": serialize_value(message.entity), + "entities": [ + { + "vectors": entity.vectors, + "entity": serialize_value(entity.entity), + } + for entity in message.entities + ], + } + +def serialize_document_embeddings(message): + return { + "metadata": { + "id": message.metadata.id, + "metadata": serialize_subgraph(message.metadata.metadata), + "user": message.metadata.user, + "collection": message.metadata.collection, + }, + "chunks": [ + { + "vectors": chunk.vectors, + "chunk": chunk.chunk.decode("utf-8"), + } + for chunk in message.chunks + ], } diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 9aea1acf..5f131c2b 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -31,6 +31,7 @@ from . subscriber import Subscriber from . text_completion import TextCompletionRequestor from . prompt import PromptRequestor from . graph_rag import GraphRagRequestor +from . document_rag import DocumentRagRequestor from . triples_query import TriplesQueryRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor from . embeddings import EmbeddingsRequestor @@ -40,11 +41,14 @@ from . dbpedia import DbpediaRequestor from . internet_search import InternetSearchRequestor from . triples_stream import TriplesStreamEndpoint from . graph_embeddings_stream import GraphEmbeddingsStreamEndpoint +from . document_embeddings_stream import DocumentEmbeddingsStreamEndpoint from . triples_load import TriplesLoadEndpoint from . graph_embeddings_load import GraphEmbeddingsLoadEndpoint +from . document_embeddings_load import DocumentEmbeddingsLoadEndpoint from . mux import MuxEndpoint from . document_load import DocumentLoadSender from . text_load import TextLoadSender +from . metrics import MetricsEndpoint from . endpoint import ServiceEndpoint from . auth import Authenticator @@ -54,6 +58,7 @@ logger.setLevel(logging.INFO) default_pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650") default_pulsar_api_key = os.getenv("PULSAR_API_KEY", None) +default_prometheus_url = os.getenv("PROMETHEUS_URL", "http://prometheus:9090") default_timeout = 600 default_port = 8088 default_api_token = os.getenv("GATEWAY_SECRET", "") @@ -72,6 +77,13 @@ class Api: self.pulsar_host = config.get("pulsar_host", default_pulsar_host) self.pulsar_api_key = config.get("pulsar_api_key", default_pulsar_api_key) + self.prometheus_url = config.get( + "prometheus_url", default_prometheus_url, + ) + + if not self.prometheus_url.endswith("/"): + self.prometheus_url += "/" + api_token = config.get("api_token", default_api_token) # Token not set, or token equal empty string means no auth @@ -93,6 +105,10 @@ class Api: pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, pulsar_api_key=self.pulsar_api_key, ), + "document-rag": DocumentRagRequestor( + pulsar_host=self.pulsar_host, timeout=self.timeout, + auth = self.auth, + ), "triples-query": TriplesQueryRequestor( pulsar_host=self.pulsar_host, timeout=self.timeout, auth = self.auth, pulsar_api_key=self.pulsar_api_key, @@ -142,6 +158,10 @@ class Api: endpoint_path = "/api/v1/graph-rag", auth=self.auth, requestor = self.services["graph-rag"], ), + ServiceEndpoint( + endpoint_path = "/api/v1/document-rag", auth=self.auth, + requestor = self.services["document-rag"], + ), ServiceEndpoint( endpoint_path = "/api/v1/triples-query", auth=self.auth, requestor = self.services["triples-query"], @@ -189,6 +209,10 @@ class Api: pulsar_api_key=self.pulsar_api_key, auth = self.auth, ), + DocumentEmbeddingsStreamEndpoint( + pulsar_host=self.pulsar_host, + auth = self.auth, + ), TriplesLoadEndpoint( pulsar_host=self.pulsar_host, auth = self.auth, @@ -199,12 +223,21 @@ class Api: pulsar_api_key=self.pulsar_api_key, auth = self.auth, ), + DocumentEmbeddingsLoadEndpoint( + pulsar_host=self.pulsar_host, + auth = self.auth, + ), MuxEndpoint( pulsar_host=self.pulsar_host, auth = self.auth, services = self.services, pulsar_api_key=self.pulsar_api_key, ), + MetricsEndpoint( + endpoint_path = "/api/v1/metrics", + prometheus_url = self.prometheus_url, + auth = self.auth, + ), ] for ep in self.endpoints: @@ -239,6 +272,12 @@ def run(): help=f'Pulsar API key', ) + parser.add_argument( + '-m', '--prometheus-url', + default=default_prometheus_url, + help=f'Prometheus URL (default: {default_prometheus_url})', + ) + parser.add_argument( '--port', type=int, diff --git a/trustgraph-flow/trustgraph/gateway/socket.py b/trustgraph-flow/trustgraph/gateway/socket.py index fd408d7b..4adc336f 100644 --- a/trustgraph-flow/trustgraph/gateway/socket.py +++ b/trustgraph-flow/trustgraph/gateway/socket.py @@ -44,7 +44,10 @@ class SocketEndpoint: return web.HTTPUnauthorized() running = Running() - ws = web.WebSocketResponse() + + # 50MB max message size + ws = web.WebSocketResponse(max_msg_size=52428800) + await ws.prepare(request) try: diff --git a/trustgraph-flow/trustgraph/gateway/subscriber.py b/trustgraph-flow/trustgraph/gateway/subscriber.py index 09622b3e..60406f6b 100644 --- a/trustgraph-flow/trustgraph/gateway/subscriber.py +++ b/trustgraph-flow/trustgraph/gateway/subscriber.py @@ -7,7 +7,7 @@ import time class Subscriber: def __init__(self, pulsar_host, topic, subscription, consumer_name, pulsar_api_key=None, - schema=None, max_size=100): + schema=None, max_size=100, listener=None): self.pulsar_host = pulsar_host self.pulsar_api_key = pulsar_api_key self.topic = topic @@ -18,6 +18,7 @@ class Subscriber: self.full = {} self.max_size = max_size self.lock = threading.Lock() + self.listener_name = listener def start(self): self.task = threading.Thread(target=self.run) @@ -32,12 +33,14 @@ class Subscriber: if self.pulsar_api_key: auth = pulsar.AuthenticationToken(self.pulsar_api_key) client = pulsar.Client( - self.pulsar_host, - authentication=auth, + self.pulsar_host, + authentication=auth, + listener_name=self.listener_name, ) else: client = pulsar.Client( - self.pulsar_host, + self.pulsar_host, + listener_name=self.listener_name, ) consumer = client.subscribe( diff --git a/trustgraph-flow/trustgraph/gateway/text_load.py b/trustgraph-flow/trustgraph/gateway/text_load.py index 562ca779..2857b601 100644 --- a/trustgraph-flow/trustgraph/gateway/text_load.py +++ b/trustgraph-flow/trustgraph/gateway/text_load.py @@ -37,7 +37,7 @@ class TextLoadSender(ServiceSender): return TextDocument( metadata=Metadata( id=body.get("id"), - metabody=metadata, + metadata=metadata, user=body.get("user", "trustgraph"), collection=body.get("collection", "default"), ), diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/triples_stream.py index 2507a34d..b98ba437 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/triples_stream.py @@ -29,6 +29,16 @@ class TriplesStreamEndpoint(SocketEndpoint): schema=JsonSchema(Triples) ) + async def listener(self, ws, running): + + worker = asyncio.create_task( + self.async_thread(ws, running) + ) + + await super(TriplesStreamEndpoint, self).listener(ws, running) + + await worker + async def start(self): self.subscriber.start() @@ -44,6 +54,9 @@ class TriplesStreamEndpoint(SocketEndpoint): resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_triples(resp)) + except TimeoutError: + continue + except queue.Empty: continue diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index 4db7dbf1..90be6962 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -158,25 +158,15 @@ class Processor(ConsumerProducer): except TooManyRequests: - print("Send rate limit response...", flush=True) + print("Rate limit...") - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index a3edb859..f5ecb8d6 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -4,10 +4,9 @@ Simple LLM service, performs text prompt completion using the Azure OpenAI endpoit service. Input is prompt, output is response. """ -import requests import json from prometheus_client import Histogram -from openai import AzureOpenAI +from openai import AzureOpenAI, RateLimitError import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error @@ -126,30 +125,27 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp.choices[0].message.content, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) - self.producer.send(r, properties={"id": id}) - - except TooManyRequests: - - print("Send rate limit response...", flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, + response=resp.choices[0].message.content, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model ) self.producer.send(r, properties={"id": id}) - self.consumer.acknowledge(msg) + except RateLimitError: + + print("Send rate limit response...", flush=True) + + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 01ce837d..5cfd8907 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -87,8 +87,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits? - with __class__.text_completion_metric.time(): response = message = self.claude.messages.create( @@ -117,34 +115,26 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) + r = TextCompletionResponse( + response=resp, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ) self.send(r, properties={"id": id}) print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: + except anthropic.RateLimitError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index d03e1554..5b8e3ba9 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -112,27 +112,15 @@ class Processor(ConsumerProducer): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except TooManyRequests: + except cohere.TooManyRequestsError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index a249998d..5d5b23a0 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -88,7 +88,8 @@ class Processor(ConsumerProducer): HarmCategory.HARM_CATEGORY_HARASSMENT: block_level, HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: block_level, HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: block_level, - # There is a documentation conflict on whether or not CIVIC_INTEGRITY is a valid category + # There is a documentation conflict on whether or not + # CIVIC_INTEGRITY is a valid category # HarmCategory.HARM_CATEGORY_CIVIC_INTEGRITY: block_level, } @@ -122,8 +123,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits? - with __class__.text_completion_metric.time(): chat_session = self.llm.start_chat( @@ -140,35 +139,30 @@ class Processor(ConsumerProducer): print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) - r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) + r = TextCompletionResponse( + response=resp, + error=None, + in_token=inputtokens, + out_token=outputtokens, + model=self.model + ) self.send(r, properties={"id": id}) print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit except ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the default handler + raise TooManyRequests() except Exception as e: - print(f"Exception: {e}") + # Apart from rate limits, treat all exceptions as unrecoverable + + print(type(e), flush=True) + print(f"Exception: {e}", flush=True) print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index 274948a8..65a2b171 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -126,26 +126,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: - - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # SLM, presumably there aren't rate limits except Exception as e: diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 00d44f6d..8c5bd3dc 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -100,26 +100,7 @@ class Processor(ConsumerProducer): print("Done.", flush=True) - # FIXME: Wrong exception, don't know what this LLM throws - # for a rate limit - except TooManyRequests: - - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # SLM, presumably no rate limits except Exception as e: diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index c874943e..c2b948d5 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -4,7 +4,7 @@ Simple LLM service, performs text prompt completion using OpenAI. Input is prompt, output is response. """ -from openai import OpenAI +from openai import OpenAI, RateLimitError from prometheus_client import Histogram import os @@ -87,8 +87,6 @@ class Processor(ConsumerProducer): try: - # FIXME: Rate limits - with __class__.text_completion_metric.time(): resp = self.openai.chat.completions.create( @@ -134,27 +132,15 @@ class Processor(ConsumerProducer): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except TooManyRequests: + except openai.RateLimitError: - print("Send rate limit response...", flush=True) - - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py index 7bb5133a..dd53862f 100755 --- a/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py @@ -30,6 +30,8 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + #optional api key + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -39,10 +41,11 @@ class Processor(ConsumerProducer): "input_schema": DocumentEmbeddingsRequest, "output_schema": DocumentEmbeddingsResponse, "store_uri": store_uri, + "api_key": api_key, } ) - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) def handle(self, msg): @@ -111,7 +114,13 @@ class Processor(ConsumerProducer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Milvus store URI (default: {default_store_uri})' + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py index c2dcaa4c..4e9492c6 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py @@ -30,6 +30,7 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -39,10 +40,11 @@ class Processor(ConsumerProducer): "input_schema": GraphEmbeddingsRequest, "output_schema": GraphEmbeddingsResponse, "store_uri": store_uri, + "api_key": api_key, } ) - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) def create_value(self, ent): if ent.startswith("http://") or ent.startswith("https://"): @@ -137,7 +139,13 @@ class Processor(ConsumerProducer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Milvus store URI (default: {default_store_uri})' + help=f'Qdrant store URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'API key for qdrant (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py index 4245784d..22fbf84d 100755 --- a/trustgraph-flow/trustgraph/query/triples/cassandra/service.py +++ b/trustgraph-flow/trustgraph/query/triples/cassandra/service.py @@ -26,6 +26,8 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -35,10 +37,14 @@ class Processor(ConsumerProducer): "input_schema": TriplesQueryRequest, "output_schema": TriplesQueryResponse, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) self.graph_host = [graph_host] + self.username = graph_username + self.password = graph_password self.table = None def create_value(self, ent): @@ -56,10 +62,17 @@ class Processor(ConsumerProducer): table = (v.user, v.collection) if table != self.table: - self.tg = TrustGraph( - hosts=self.graph_host, - keyspace=v.user, table=v.collection, - ) + if self.username and self.password: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.user, table=v.collection, + username=self.username, password=self.password + ) + else: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.user, table=v.collection, + ) self.table = table # Sender-produced ID @@ -176,6 +189,19 @@ class Processor(ConsumerProducer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) + def run(): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py index 00f9d5b5..bfa6c123 100755 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py @@ -3,15 +3,16 @@ Accepts entity/vector pairs and writes them to a Milvus store. """ -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue -from .... log_level import LogLevel from .... direct.milvus_doc_embeddings import DocVectors + +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue +from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_store_uri = 'http://localhost:19530' @@ -27,7 +28,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "store_uri": store_uri, } ) @@ -38,11 +39,16 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if v.chunk != "" and v.chunk is not None: - for vec in v.vectors: - self.vecstore.insert(vec, chunk) + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue + + for vec in emb.vectors: + + if chunk != "" and v.chunk is not None: + for vec in v.vectors: + self.vecstore.insert(vec, chunk) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py index 24cfcb78..c59ecd7b 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py @@ -11,14 +11,14 @@ import time import uuid import os -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_api_key = os.getenv("PINECONE_API_KEY", "not-specified") default_cloud = "aws" @@ -54,7 +54,7 @@ class Processor(Consumer): **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "url": self.url, } ) @@ -65,71 +65,74 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "" or chunk is None: continue - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d-" + v.metadata.user + "-" + str(dim) - ) + for vec in v.vectors: - if index_name != self.last_index_name: + dim = len(vec) + collection = ( + "d-" + v.metadata.user + "-" + str(dim) + ) - if not self.pinecone.has_index(index_name): + if index_name != self.last_index_name: - try: + if not self.pinecone.has_index(index_name): - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + for i in range(0, 1000): - time.sleep(1) + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + time.sleep(1) - except Exception as e: - print("Pinecone index creation failed") - raise e + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) - print(f"Index {index_name} created", flush=True) + except Exception as e: + print("Pinecone index creation failed") + raise e - self.last_index_name = index_name + print(f"Index {index_name} created", flush=True) - index = self.pinecone.Index(index_name) + self.last_index_name = index_name - records = [ - { - "id": id, - "values": vec, - "metadata": { "doc": chunk }, - } - ] + index = self.pinecone.Index(index_name) - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + records = [ + { + "id": id, + "values": vec, + "metadata": { "doc": chunk }, + } + ] + + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py index 813c4f29..a01fd9e4 100644 --- a/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py @@ -8,14 +8,14 @@ from qdrant_client.models import PointStruct from qdrant_client.models import Distance, VectorParams import uuid -from .... schema import ChunkEmbeddings -from .... schema import chunk_embeddings_ingest_queue +from .... schema import DocumentEmbeddings +from .... schema import document_embeddings_store_queue from .... log_level import LogLevel from .... base import Consumer module = ".".join(__name__.split(".")[1:-1]) -default_input_queue = chunk_embeddings_ingest_queue +default_input_queue = document_embeddings_store_queue default_subscriber = module default_store_uri = 'http://localhost:6333' @@ -26,13 +26,15 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { "input_queue": input_queue, "subscriber": subscriber, - "input_schema": ChunkEmbeddings, + "input_schema": DocumentEmbeddings, "store_uri": store_uri, + "api_key": api_key, } ) @@ -44,47 +46,48 @@ class Processor(Consumer): v = msg.value() - chunk = v.chunk.decode("utf-8") + for emb in v.chunks: - if chunk == "": return + chunk = emb.chunk.decode("utf-8") + if chunk == "": return - for vec in v.vectors: + for vec in emb.vectors: - dim = len(vec) - collection = ( - "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + - str(dim) - ) + dim = len(vec) + collection = ( + "d_" + v.metadata.user + "_" + v.metadata.collection + "_" + + str(dim) + ) - if collection != self.last_collection: + if collection != self.last_collection: - if not self.client.collection_exists(collection): + if not self.client.collection_exists(collection): - try: - self.client.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.DOT - ), + try: + self.client.create_collection( + collection_name=collection, + vectors_config=VectorParams( + size=dim, distance=Distance.COSINE + ), + ) + except Exception as e: + print("Qdrant collection creation failed") + raise e + + self.last_collection = collection + + self.client.upsert( + collection_name=collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "doc": chunk, + } ) - except Exception as e: - print("Qdrant collection creation failed") - raise e - - self.last_collection = collection - - self.client.upsert( - collection_name=collection, - points=[ - PointStruct( - id=str(uuid.uuid4()), - vector=vec, - payload={ - "doc": chunk, - } - ) - ] - ) + ] + ) @staticmethod def add_args(parser): @@ -96,7 +99,13 @@ class Processor(Consumer): parser.add_argument( '-t', '--store-uri', default=default_store_uri, - help=f'Qdrant store URI (default: {default_store_uri})' + help=f'Qdrant URI (default: {default_store_uri})' + ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key (default: None)' ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py index 98fe7915..e1379577 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py @@ -38,9 +38,11 @@ class Processor(Consumer): v = msg.value() - if v.entity.value != "": - for vec in v.vectors: - self.vecstore.insert(vec, v.entity.value) + for entity in v.entities: + + if entity.entity.value != "" and entity.entity.value is not None: + for vec in entity.vectors: + self.vecstore.insert(vec, entity.entity.value) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index b918c10b..a32ff627 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -60,76 +60,83 @@ class Processor(Consumer): self.last_index_name = None + def create_index(self, index_name, dim): + + self.pinecone.create_index( + name = index_name, + dimension = dim, + metric = "cosine", + spec = ServerlessSpec( + cloud = self.cloud, + region = self.region, + ) + ) + + for i in range(0, 1000): + + if self.pinecone.describe_index( + index_name + ).status["ready"]: + break + + time.sleep(1) + + if not self.pinecone.describe_index( + index_name + ).status["ready"]: + raise RuntimeError( + "Gave up waiting for index creation" + ) + def handle(self, msg): v = msg.value() id = str(uuid.uuid4()) - if v.entity.value == "" or v.entity.value is None: return + for entity in v.entities: - for vec in v.vectors: + if entity.entity.value == "" or entity.entity.value is None: + continue - dim = len(vec) + for vec in entity.vectors: - index_name = ( - "t-" + v.metadata.user + "-" + str(dim) - ) + dim = len(vec) - if index_name != self.last_index_name: + index_name = ( + "t-" + v.metadata.user + "-" + str(dim) + ) - if not self.pinecone.has_index(index_name): + if index_name != self.last_index_name: - try: + if not self.pinecone.has_index(index_name): - self.pinecone.create_index( - name = index_name, - dimension = dim, - metric = "cosine", - spec = ServerlessSpec( - cloud = self.cloud, - region = self.region, - ) - ) + try: - for i in range(0, 1000): + self.create_index(index_name, dim) - if self.pinecone.describe_index( - index_name - ).status["ready"]: - break + except Exception as e: + print("Pinecone index creation failed") + raise e - time.sleep(1) + print(f"Index {index_name} created", flush=True) - if not self.pinecone.describe_index( - index_name - ).status["ready"]: - raise RuntimeError( - "Gave up waiting for index creation" - ) + self.last_index_name = index_name - except Exception as e: - print("Pinecone index creation failed") - raise e + index = self.pinecone.Index(index_name) - print(f"Index {index_name} created", flush=True) + records = [ + { + "id": id, + "values": vec, + "metadata": { "entity": entity.entity.value }, + } + ] - self.last_index_name = index_name - - index = self.pinecone.Index(index_name) - - records = [ - { - "id": id, - "values": vec, - "metadata": { "entity": v.entity.value }, - } - ] - - index.upsert( - vectors = records, - namespace = v.metadata.collection, - ) + index.upsert( + vectors = records, + namespace = v.metadata.collection, + ) @staticmethod def add_args(parser): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py index 47b53979..ffe33565 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py @@ -26,6 +26,7 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) store_uri = params.get("store_uri", default_store_uri) + api_key = params.get("api_key", None) super(Processor, self).__init__( **params | { @@ -33,56 +34,67 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": GraphEmbeddings, "store_uri": store_uri, + "api_key": api_key, } ) self.last_collection = None - self.client = QdrantClient(url=store_uri) + self.client = QdrantClient(url=store_uri, api_key=api_key) + + def get_collection(self, dim, user, collection): + + cname = ( + "t_" + user + "_" + collection + "_" + str(dim) + ) + + if cname != self.last_collection: + + if not self.client.collection_exists(cname): + + try: + self.client.create_collection( + collection_name=cname, + vectors_config=VectorParams( + size=dim, distance=Distance.COSINE + ), + ) + except Exception as e: + print("Qdrant collection creation failed") + raise e + + self.last_collection = cname + + return cname def handle(self, msg): v = msg.value() - if v.entity.value == "" or v.entity.value is None: return + for entity in v.entities: - for vec in v.vectors: + if entity.entity.value == "" or entity.entity.value is None: return - dim = len(vec) - collection = ( - "t_" + v.metadata.user + "_" + v.metadata.collection + "_" + - str(dim) - ) + for vec in entity.vectors: - if collection != self.last_collection: + dim = len(vec) - if not self.client.collection_exists(collection): + collection = self.get_collection( + dim, v.metadata.user, v.metadata.collection + ) - try: - self.client.create_collection( - collection_name=collection, - vectors_config=VectorParams( - size=dim, distance=Distance.COSINE - ), + self.client.upsert( + collection_name=collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "entity": entity.entity.value, + } ) - except Exception as e: - print("Qdrant collection creation failed") - raise e - - self.last_collection = collection - - self.client.upsert( - collection_name=collection, - points=[ - PointStruct( - id=str(uuid.uuid4()), - vector=vec, - payload={ - "entity": v.entity.value, - } - ) - ] - ) + ] + ) @staticmethod def add_args(parser): @@ -96,6 +108,12 @@ class Processor(Consumer): default=default_store_uri, help=f'Qdrant store URI (default: {default_store_uri})' ) + + parser.add_argument( + '-k', '--api-key', + default=None, + help=f'Qdrant API key' + ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py index d44864fe..fc8f6686 100755 --- a/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/rows/cassandra/write.py @@ -29,6 +29,8 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -36,10 +38,16 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": Rows, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) - - self.cluster = Cluster(graph_host.split(",")) + + if graph_username and graph_password: + auth_provider = PlainTextAuthProvider(username=graph_username, password=graph_password) + self.cluster = Cluster(graph_host.split(","), auth_provider=auth_provider) + else: + self.cluster = Cluster(graph_host.split(",")) self.session = self.cluster.connect() self.tables = set() @@ -120,6 +128,18 @@ class Processor(Consumer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py index e7078e08..d940d0ec 100755 --- a/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/cassandra/write.py @@ -28,6 +28,8 @@ class Processor(Consumer): input_queue = params.get("input_queue", default_input_queue) subscriber = params.get("subscriber", default_subscriber) graph_host = params.get("graph_host", default_graph_host) + graph_username = params.get("graph_username", None) + graph_password = params.get("graph_password", None) super(Processor, self).__init__( **params | { @@ -35,10 +37,14 @@ class Processor(Consumer): "subscriber": subscriber, "input_schema": Triples, "graph_host": graph_host, + "graph_username": graph_username, + "graph_password": graph_password, } ) - + self.graph_host = [graph_host] + self.username = graph_username + self.password = graph_password self.table = None def handle(self, msg): @@ -52,10 +58,17 @@ class Processor(Consumer): self.tg = None try: - self.tg = TrustGraph( - hosts=self.graph_host, - keyspace=v.metadata.user, table=v.metadata.collection, - ) + if self.username and self.password: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.metadata.user, table=v.metadata.collection, + username=self.username, password=self.password + ) + else: + self.tg = TrustGraph( + hosts=self.graph_host, + keyspace=v.metadata.user, table=v.metadata.collection, + ) except Exception as e: print("Exception", e, flush=True) time.sleep(1) @@ -82,6 +95,18 @@ class Processor(Consumer): default="localhost", help=f'Graph host (default: localhost)' ) + + parser.add_argument( + '--graph-username', + default=None, + help=f'Cassandra username' + ) + + parser.add_argument( + '--graph-password', + default=None, + help=f'Cassandra password' + ) def run(): diff --git a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py index f106170a..7295e691 100755 --- a/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/memgraph/write.py @@ -55,6 +55,14 @@ class Processor(Consumer): def create_indexes(self, session): + # Race condition, index creation failure is ignored. Right thing + # to do if the index already exists. Wrong thing to do if it's + # because the store is not up yet + + # In real-world cases, Memgraph will start up quicker than Pulsar + # and this process will restart several times until Pulsar arrives, + # so should be safe + print("Create indexes...", flush=True) try: diff --git a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py index 1aa25aa8..18b40129 100755 --- a/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py +++ b/trustgraph-flow/trustgraph/storage/triples/neo4j/write.py @@ -50,6 +50,50 @@ class Processor(Consumer): self.io = GraphDatabase.driver(graph_host, auth=(username, password)) + with self.io.session(database=self.db) as session: + self.create_indexes(session) + + def create_indexes(self, session): + + # Race condition, index creation failure is ignored. Right thing + # to do if the index already exists. Wrong thing to do if it's + # because the store is not up yet + + # In real-world cases, Neo4j will start up quicker than Pulsar + # and this process will restart several times until Pulsar arrives, + # so should be safe + + print("Create indexes...", flush=True) + + try: + session.run( + "CREATE INDEX Node_uri FOR (n:Node) ON (n.uri)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + try: + session.run( + "CREATE INDEX Literal_value FOR (n:Literal) ON (n.value)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + try: + session.run( + "CREATE INDEX Rel_uri FOR ()-[r:Rel]-() ON (r.uri)", + ) + except Exception as e: + print(e, flush=True) + # Maybe index already exists + print("Index create failure ignored", flush=True) + + print("Index creation done", flush=True) + def create_node(self, uri): print("Create node", uri) diff --git a/trustgraph-vertexai/setup.py b/trustgraph-vertexai/setup.py index 7f9c2923..1258fea9 100644 --- a/trustgraph-vertexai/setup.py +++ b/trustgraph-vertexai/setup.py @@ -34,7 +34,7 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", + "trustgraph-base>=0.21,<0.22", "pulsar-client", "google-cloud-aiplatform", "prometheus-client", diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index cb817836..d6a2efec 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -178,25 +178,15 @@ class Processor(ConsumerProducer): except google.api_core.exceptions.ResourceExhausted as e: - print("Send rate limit response...", flush=True) + print("Hit rate limit:", e, flush=True) - r = TextCompletionResponse( - error=Error( - type = "rate-limit", - message = str(e), - ), - response=None, - in_token=None, - out_token=None, - model=None, - ) - - self.producer.send(r, properties={"id": id}) - - self.consumer.acknowledge(msg) + # Leave rate limit retries to the base handler + raise TooManyRequests() except Exception as e: + # Apart from rate limits, treat all exceptions as unrecoverable + print(f"Exception: {e}") print("Send error response...", flush=True) diff --git a/trustgraph/setup.py b/trustgraph/setup.py index a964ff06..d7185e66 100644 --- a/trustgraph/setup.py +++ b/trustgraph/setup.py @@ -34,12 +34,12 @@ setuptools.setup( python_requires='>=3.8', download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz", install_requires=[ - "trustgraph-base>=0.18,<0.19", - "trustgraph-bedrock>=0.18,<0.19", - "trustgraph-cli>=0.18,<0.19", - "trustgraph-embeddings-hf>=0.18,<0.19", - "trustgraph-flow>=0.18,<0.19", - "trustgraph-vertexai>=0.18,<0.19", + "trustgraph-base>=0.21,<0.22", + "trustgraph-bedrock>=0.21,<0.22", + "trustgraph-cli>=0.21,<0.22", + "trustgraph-embeddings-hf>=0.21,<0.22", + "trustgraph-flow>=0.21,<0.22", + "trustgraph-vertexai>=0.21,<0.22", ], scripts=[ ]