apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-05 19:32:11 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Tried to add Cohere support

Browse source
This commit is contained in:
JackColquitt 2024-07-31 15:42:06 -07:00
parent 315b4d81d6
commit db006a4b7b
24 changed files with 19258 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
.DS_Store vendored Normal file
View file

Binary file not shown.

102
.ipynb_checkpoints/README-checkpoint.md Normal file
View file

@ -0,0 +1,102 @@
# TrustGraph
## Introduction
TrustGraph is a true end-to-end (e2e) knowledge pipeline that performs a `naive extraction` on a text corpus
to build a RDF style knowledge graph coupled with a `RAG` service compatible with cloud LLMs and open-source
SLMs (Small Language Models).
The pipeline processing components are interconnected with a pub/sub engine to
maximize modularity and enable new knowledge processing functions. The core processing components decode documents,
chunk text, perform embeddings, apply a local SLM/LLM, call a LLM API, and generate LM predictions.
The processing showcases the reliability and efficiences of Graph RAG algorithms which can capture
contextual language flags that are missed in conventional RAG approaches. Graph querying algorithms enable retrieving
not just relevant knowledge but language cues essential to understanding semantic uses unique to a text corpus.
Processing modules are executed in containers. Processing can be scaled-up
by deploying multiple containers.
### Features
- PDF decoding
- Text chunking
- Inference of LMs deployed with [Ollama](https://ollama.com)
- Inference of LLMs: Claude, VertexAI and AzureAI serverless endpoints
- Application of a [HuggingFace](https://hf.co) embeddings models
- [RDF](https://www.w3.org/TR/rdf12-schema/)-aligned Knowledge Graph extraction
- Graph edge loading into [Apache Cassandra](https://github.com/apache/cassandra)
- Storing embeddings in [Milvus](https://github.com/milvus-io/milvus)
- Embedding query service
- Graph RAG query service
- All procesing integrates with [Apache Pulsar](https://github.com/apache/pulsar/)
- Containers, so can be deployed using Docker Compose or Kubernetes
- Plug'n'play architecture: switch different LLM modules to suit your needs
## Architecture
![architecture](architecture.png)
TrustGraph is designed to be modular to support as many Language Models and environments as possible. A natural
fit for a modular architecture is to decompose functions into a set modules connected through a pub/sub backbone.
[Apache Pulsar](https://github.com/apache/pulsar/) serves as this pub/sub backbone. Pulsar acts as the data broker
managing inputs and outputs between modules.
**Pulsar Workflows**:
- For processing flows, Pulsar accepts the output of a processing module
and queues it for input to the next subscribed module.
- For services such as LLMs and embeddings, Pulsar provides a client/server
model. A Pulsar queue is used as the input to the service. When
processed, the output is then delivered to a separate queue where a client
subscriber can request that output.
The entire architecture, the pub/sub backbone and set of modules, is bundled into a single Python package. A container image with the
package installed can also run the entire architecture.
## Core Modules
- `chunker-recursive` - Accepts text documents and uses LangChain recursive
chunking algorithm to produce smaller text chunks.
- `embeddings-hf` - A service which analyses text and returns a vector
embedding using one of the HuggingFace embeddings models.
- `embeddings-vectorize` - Uses an embeddings service to get a vector
embedding which is added to the processor payload.
- `graph-rag` - A query service which applies a Graph RAG algorithm to
provide a response to a text prompt.
- `graph-write-cassandra` - Takes knowledge graph edges and writes them to
a Cassandra store.
- `kg-extract-definitions` - knowledge extractor - examines text and
produces graph edges.
describing discovered terms and also their defintions. Definitions are
derived using the input documents.
- `kg-extract-relationships` - knowledge extractor - examines text and
produces graph edges describing the relationships between discovered
terms.
- `loader` - Takes a document and loads into the processing pipeline. Used
e.g. to add PDF documents.
- `pdf-decoder` - Takes a PDF doc and emits text extracted from the document.
Text extraction from PDF is not a perfect science as PDF is a printable
format. For instance, the wrapping of text between lines in a PDF document
is not semantically encoded, so the decoder will see wrapped lines as
space-separated.
- `vector-write-milvus` - Takes vector-entity mappings and records them
in the vector embeddings store.
## LM Specific Modules
- `llm-azure-text` - Sends request to AzureAI serverless endpoint
- `llm-claude-text` - Sends request to Anthropic's API
- `llm-ollama-text` - Sends request to LM running using Ollama
- `llm-vertexai-text` - Sends request to model available through VertexAI API
## Quickstart Guide
See [Quickstart on Docker Compose](docs/README.quickstart-docker-compose.md)
## Development Guide
See [Development on trustgraph](docs/README.development.md)

179
.ipynb_checkpoints/docker-compose-azure-checkpoint.yaml Normal file
View file

@ -0,0 +1,179 @@
volumes:
cassandra:
pulsar-conf:
pulsar-data:
etcd:
minio-data:
milvus:
services:
cassandra:
image: docker.io/cassandra:4.1.5
ports:
- "9042:9042"
volumes:
- "cassandra:/var/lib/cassandra"
restart: on-failure:100
pulsar:
image: docker.io/apachepulsar/pulsar:3.3.0
command: bin/pulsar standalone
ports:
- "6650:6650"
- "8080:8080"
volumes:
- "pulsar-conf:/pulsar/conf"
- "pulsar-data:/pulsar/data"
restart: on-failure:100
pulsar-manager:
image: docker.io/apachepulsar/pulsar-manager:v0.3.0
ports:
- "9527:9527"
- "7750:7750"
environment:
SPRING_CONFIGURATION_FILE: /pulsar-manager/pulsar-manager/application.properties
restart: on-failure:100
etcd:
image: quay.io/coreos/etcd:v3.5.5
command:
- "etcd"
- "-advertise-client-urls=http://127.0.0.1:2379"
- "-listen-client-urls"
- "http://0.0.0.0:2379"
- "--data-dir"
- "/etcd"
environment:
ETCD_AUTO_COMPACTION_MODE: revision
ETCD_AUTO_COMPACTION_RETENTION: "1000"
ETCD_QUOTA_BACKEND_BYTES: "4294967296"
ETCD_SNAPSHOT_COUNT: "50000"
ports:
- "2379:2379"
volumes:
- "etcd:/etcd"
restart: on-failure:100
minio:
image: docker.io/minio/minio:RELEASE.2024-07-04T14-25-45Z
command:
- "minio"
- "server"
- "/minio_data"
- "--console-address"
- ":9001"
environment:
MINIO_ROOT_USER: minioadmin
MINIO_ROOT_PASSWORD: minioadmin
ports:
- "9001:9001"
volumes:
- "minio-data:/minio_data"
restart: on-failure:100
milvus:
image: docker.io/milvusdb/milvus:v2.4.5
command:
- "milvus"
- "run"
- "standalone"
environment:
ETCD_ENDPOINTS: etcd:2379
MINIO_ADDRESS: minio:9000
ports:
- "9091:9091"
- "19530:19530"
volumes:
- "milvus:/var/lib/milvus"
restart: on-failure:100
pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "pdf-decoder"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
chunker:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "chunker-recursive"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "embeddings-vectorize"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "embeddings-hf"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "kg-extract-definitions"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "kg-extract-relationships"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
vector-write:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "vector-write-milvus"
- "-p"
- "pulsar://pulsar:6650"
- "-t"
- "http://milvus:19530"
restart: on-failure:100
graph-write:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "graph-write-cassandra"
- "-p"
- "pulsar://pulsar:6650"
- "-g"
- "cassandra"
restart: on-failure:100
llm:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "llm-azure-text"
- "-p"
- "pulsar://pulsar:6650"
- "-k"
- ${AZURE_TOKEN}
- "-e"
- ${AZURE_ENDPOINT}
restart: on-failure:100
graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.1.16
command:
- "graph-rag"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100

219
docker-compose-cohere.yaml Normal file
View file

@ -0,0 +1,219 @@
volumes:
cassandra:
pulsar-conf:
pulsar-data:
etcd:
minio-data:
milvus:
prometheus-data:
grafana-storage:
services:
cassandra:
image: docker.io/cassandra:4.1.5
ports:
- "9042:9042"
volumes:
- "cassandra:/var/lib/cassandra"
restart: on-failure:100
pulsar:
image: docker.io/apachepulsar/pulsar:3.3.0
command: bin/pulsar standalone
ports:
- "6650:6650"
- "8080:8080"
volumes:
- "pulsar-conf:/pulsar/conf"
- "pulsar-data:/pulsar/data"
restart: on-failure:100
init-pulsar:
image: docker.io/apachepulsar/pulsar:3.3.0
command:
- "sh"
- "-c"
- "pulsar-admin --admin-url http://pulsar:8080 tenants create tg && pulsar-admin --admin-url http://pulsar:8080 namespaces create tg/flow && pulsar-admin --admin-url http://pulsar:8080 namespaces create tg/request && pulsar-admin --admin-url http://pulsar:8080 namespaces create tg/response && pulsar-admin --admin-url http://pulsar:8080 namespaces set-retention --size -1 --time 3m tg/response"
depends_on:
pulsar:
condition: service_started
restart: on-failure:100
pulsar-manager:
image: docker.io/apachepulsar/pulsar-manager:v0.3.0
ports:
- "9527:9527"
- "7750:7750"
environment:
SPRING_CONFIGURATION_FILE: /pulsar-manager/pulsar-manager/application.properties
restart: on-failure:100
etcd:
image: quay.io/coreos/etcd:v3.5.5
command:
- "etcd"
- "-advertise-client-urls=http://127.0.0.1:2379"
- "-listen-client-urls"
- "http://0.0.0.0:2379"
- "--data-dir"
- "/etcd"
environment:
ETCD_AUTO_COMPACTION_MODE: revision
ETCD_AUTO_COMPACTION_RETENTION: "1000"
ETCD_QUOTA_BACKEND_BYTES: "4294967296"
ETCD_SNAPSHOT_COUNT: "50000"
ports:
- "2379:2379"
volumes:
- "etcd:/etcd"
restart: on-failure:100
minio:
image: docker.io/minio/minio:RELEASE.2024-07-04T14-25-45Z
command:
- "minio"
- "server"
- "/minio_data"
- "--console-address"
- ":9001"
environment:
MINIO_ROOT_USER: minioadmin
MINIO_ROOT_PASSWORD: minioadmin
ports:
- "9001:9001"
volumes:
- "minio-data:/minio_data"
restart: on-failure:100
milvus:
image: docker.io/milvusdb/milvus:v2.4.5
command:
- "milvus"
- "run"
- "standalone"
environment:
ETCD_ENDPOINTS: etcd:2379
MINIO_ADDRESS: minio:9000
ports:
- "9091:9091"
- "19530:19530"
volumes:
- "milvus:/var/lib/milvus"
restart: on-failure:100
prometheus:
image: docker.io/prom/prometheus:v2.53.1
ports:
- "9090:9090"
volumes:
- "./prometheus:/etc/prometheus"
- "prometheus-data:/prometheus"
restart: on-failure:100
grafana:
image: docker.io/grafana/grafana:10.0.0
ports:
- "3000:3000"
volumes:
- "grafana-storage:/var/lib/grafana"
- "./grafana/dashboard.yml:/etc/grafana/provisioning/dashboards/dashboard.yml"
- "./grafana/datasource.yml:/etc/grafana/provisioning/datasources/datasource.yml"
- "./grafana/dashboard.json:/var/lib/grafana/dashboards/dashboard.json"
environment:
# GF_AUTH_ANONYMOUS_ORG_ROLE: Admin
# GF_AUTH_ANONYMOUS_ENABLED: true
# GF_ORG_ROLE: Admin
GF_ORG_NAME: trustgraph.ai
# GF_SERVER_ROOT_URL: https://example.com
restart: on-failure:100
pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "pdf-decoder"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
chunker:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "chunker-recursive"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "embeddings-vectorize"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "embeddings-hf"
- "-p"
- "pulsar://pulsar:6650"
# - "-m"
# - "mixedbread-ai/mxbai-embed-large-v1"
restart: on-failure:100
kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "kg-extract-definitions"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "kg-extract-relationships"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100
store-graph-embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "ge-write-milvus"
- "-p"
- "pulsar://pulsar:6650"
- "-t"
- "http://milvus:19530"
restart: on-failure:100
store-triples:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "triples-write-cassandra"
- "-p"
- "pulsar://pulsar:6650"
- "-g"
- "cassandra"
restart: on-failure:100
text-completion:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "text-completion-cohere"
- "-p"
- "pulsar://pulsar:6650"
- "-k"
- ${COHERE_KEY}
restart: on-failure:100
graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.5.5
command:
- "graph-rag"
- "-p"
- "pulsar://pulsar:6650"
restart: on-failure:100

320
docs/.ipynb_checkpoints/README.quickstart-docker-compose-checkpoint.md Normal file
View file

@ -0,0 +1,320 @@
## Getting Started
The `Docker Compose` files have been tested on `Linux` and `MacOS`. There are currently
no plans for `Windows` support in the immediate future.
There are 4 `Docker Compose` files depending on the desired LM deployment:
- `VertexAI` through Google Cloud
- `Claude` through Anthropic's API
- `AzureAI` serverless endpoint
- Local LM deployment through `Ollama`
Docker Compose enables the following functions:
- Run the required components for full e2e `Graph RAG` knowledge pipeline
- Check processing logs
- Load test text corpus and begin knowledge extraction
- Verify extracted graph edges and number of edges
- Run a query against the vector and graph stores to generate a response
using the chosen LM
### Clone the Repo
```
git clone https://github.com/trustgraph-ai/trustgraph trustgraph
cd trustgraph
```
### Install requirements
```
python3 -m venv env
. env/bin/activate
pip3 install pulsar-client
pip3 install cassandra-driver
export PYTHON_PATH=.
```
### Docker Compose files
Depending on your desired LM deployment, you will choose from one of the
following `Docker Compose` files:
- `docker-compose-azure.yaml`: AzureAI endpoint. Set `AZURE_TOKEN` to the secret token and
`AZURE_ENDPOINT` to the URL endpoint address for the deployed model.
- `docker-compose-claude.yaml`: Anthropic's API. Set `CLAUDE_KEY` to your API key.
- `docker-compose-ollama.yaml`: Local LM (currently using [Gemma2](https://ollama.com/library/gemma2) deployed through Ollama. Set `OLLAMA_HOST` to the machine running Ollama (e.g. `localhost` for Ollama running locally on your machine)
- `docker-compose-vertexai.yaml`: VertexAI API. Requires a `private.json` authentication file to authenticate with your GCP project. Filed should stored be at path `vertexai/private.json`.
**NOTE**: All tokens, paths, and authentication files must be set **PRIOR** to launching a `Docker Compose` file.
#### AzureAI Serverless Model Deployment
```
export AZURE_ENDPOINT=https://ENDPOINT.HOST.GOES.HERE/
export AZURE_TOKEN=TOKEN-GOES-HERE
docker-compose -f docker-compose-azure.yaml up -d
```
#### Claude through Anthropic API
```
export CLAUDE_KEY=TOKEN-GOES-HERE
docker-compose -f docker-compose-claude.yaml up -d
```
#### Ollama Hosted Model Deployment
```
export OLLAMA_HOST=localhost # Set to hostname of Ollama host
docker-compose -f docker-compose-ollama.yaml up -d
```
#### VertexAI through GCP
```
mkdir -p vertexai
cp {whatever} vertexai/private.json
docker-compose -f docker-compose-vertexai.yaml up -d
```
If you're running `SELinux` on Linux you may need to set the permissions on the
VertexAI directory so that the key file can be mounted on a Docker container using
the following command:
```
chcon -Rt svirt_sandbox_file_t vertexai/
```
### Verify Docker Containers
On first running a `Docker Compose` file, it may take a while (depending on your network connection) to pull all the necessary components. Once all of the components have been pulled, check that the TrustGraph containers are running:
```
docker ps
```
Any containers that have exited unexpectedly can be found by checking the `STATUS` field
using the following:
```
docker ps -a
```
### Warm-Up
Before proceeding, allow the system to enter a stable a working state. In general
`30 seconds` should be enough time for Pulsar to stablize.
The system uses Cassandra for a Graph store. Cassandra can take `60-70 seconds`
to achieve a working state.
### Load a Text Corpus
Create a sources directory and get a test PDF file. To demonstrate the power of TrustGraph, we're using a PDF of the public [Roger's Commision Report](https://sma.nasa.gov/SignificantIncidents/assets/rogers_commission_report.pdf) from the NASA Challenger disaster. This PDF includes complex formatting, unique terms, complex concepts, unique concepts, and information not commonly found in public knowledge sources.
```
mkdir sources
curl -o sources/Challenger-Report-Vol1.pdf https://sma.nasa.gov/SignificantIncidents/assets/rogers_commission_report.pdf
```
Load the file for knowledge extraction:
```
scripts/loader -f sources/Challenger-Report-Vol1.pdf
```
`File loaded.` indicates the PDF has been sucessfully loaded to the processing queues and extraction will begin.
### Processing Logs
At this point, many processing services are running concurrently. You can check the status of these processes with the following logs:
`PDF Decoder`:
```
docker logs trustgraph-pdf-decoder-1
```
Output should look:
```
Decoding 1f7b7055...
Done.
```
`Chunker`:
```
docker logs trustgraph-chunker-1
```
The output should be similiar to the output of the `Decode`, except it should be a sequence of many entries.
`Vectorizer`:
```
docker logs trustgraph-vectorize-1
```
Similar output to above processes, except many entries instead.
`Language Model Inference`:
```
docker logs trustgraph-llm-1
```
Output should be a sequence of entries:
```
Handling prompt fa1b98ae-70ef-452b-bcbe-21a867c5e8e2...
Send response...
Done.
```
`Knowledge Graph Definitions`:
```
docker logs trustgraph-kg-extract-definitions-1
```
Output should be an array of JSON objects with keys `entity` and `definition`:
```
Indexing 1f7b7055-p11-c1...
[
{
"entity": "Orbiter",
"definition": "A spacecraft designed for spaceflight."
},
{
"entity": "flight deck",
"definition": "The top level of the crew compartment, typically where flight controls are located."
},
{
"entity": "middeck",
"definition": "The lower level of the crew compartment, used for sleeping, working, and storing equipment."
}
]
Done.
```
`Knowledge Graph Relationshps`:
```
docker logs trustgraph-kg-extract-relationships-1
```
Output should be an array of JSON objects with keys `subject`, `predicate`, `object`, and `object-entity`:
```
Indexing 1f7b7055-p11-c3...
[
{
"subject": "Space Shuttle",
"predicate": "carry",
"object": "16 tons of cargo",
"object-entity": false
},
{
"subject": "friction",
"predicate": "generated by",
"object": "atmosphere",
"object-entity": true
}
]
Done.
```
### Graph Parsing
To check that the knowledge graph is successfully parsing data:
```
scripts/graph-show
```
The output should be a set of semantic triples in [N-Triples](https://www.w3.org/TR/rdf12-n-triples/) format.
```
http://trustgraph.ai/e/enterprise http://trustgraph.ai/e/was-carried to altitude and released for a gliding approach and landing at the Mojave Desert test center.
http://trustgraph.ai/e/enterprise http://www.w3.org/2000/01/rdf-schema#label Enterprise.
http://trustgraph.ai/e/enterprise http://www.w3.org/2004/02/skos/core#definition A prototype space shuttle orbiter used for atmospheric flight testing.
```
### Number of Graph Edges
N-Triples format is not particularly human readable. It's more useful to know how many graph edges have successfully been extracted from the text corpus:
```
scripts/graph-show | wc -l
```
The Challenger report has a long introduction with quite a bit of adminstrative text commonly found in official reports. The first few hundred graph edges mostly capture this document formatting knowledge. To fully test the ability to extract complex knowledge, wait until at least `1000` graph edges have been extracted. The full extraction for this PDF will extract many thousand graph edges.
### RAG Test Script
```
tests/test-graph-rag
```
This script forms a LM prompt asking for 20 facts regarding the Challenger disaster. Depending on how many graph edges have been extracted, the response will be similar to:
```
Here are 20 facts from the provided knowledge graph about the Space Shuttle disaster:
1. **Space Shuttle Challenger was a Space Shuttle spacecraft.**
2. **The third Spacelab mission was carried by Orbiter Challenger.**
3. **Francis R. Scobee was the Commander of the Challenger crew.**
4. **Earth-to-orbit systems are designed to transport payloads and humans from Earth's surface into orbit.**
5. **The Space Shuttle program involved the Space Shuttle.**
6. **Orbiter Challenger flew on mission 41-B.**
7. **Orbiter Challenger was used on STS-7 and STS-8 missions.**
8. **Columbia completed the orbital test.**
9. **The Space Shuttle flew 24 successful missions.**
10. **One possibility for the Space Shuttle was a winged but unmanned recoverable liquid-fuel vehicle based on the Saturn 5 rocket.**
11. **A Commission was established to investigate the space shuttle Challenger accident.**
12. **Judit h Arlene Resnik was Mission Specialist Two.**
13. **Mission 51-L was originally scheduled for December 1985 but was delayed until January 1986.**
14. **The Corporation's Space Transportation Systems Division was responsible for the design and development of the Space Shuttle Orbiter.**
15. **Michael John Smith was the Pilot of the Challenger crew.**
16. **The Space Shuttle is composed of two recoverable Solid Rocket Boosters.**
17. **The Space Shuttle provides for the broadest possible spectrum of civil/military missions.**
18. **Mission 51-L consisted of placing one satellite in orbit, deploying and retrieving Spartan, and conducting six experiments.**
19. **The Space Shuttle became the focus of NASA's near-term future.**
20. **The Commission focused its attention on safety aspects of future flights.**
```
For any errors with the `RAG` proces, check the following log:
```
docker logs -f trustgraph-graph-rag-1
```
### More RAG Test Queries
If you want to try different RAG queries, modify the `query` in the [test script](https://github.com/trustgraph-ai/trustgraph/blob/master/tests/test-graph-rag).
### Shutting Down
When shutting down the pipeline, it's best to shut down all Docker containers and volumes. Run the `docker compose down` command that corresponds to your model deployment:
`AzureAI Endpoint`:
```
docker-compose -f docker-compose-azure.yaml down --volumes
```
`Anthropic API`:
```
docker-compose -f docker-compose-claude.yaml down --volumes
```
`Ollama`:
```
docker-compose -f docker-compose-ollama.yaml down --volumes
```
`VertexAI API`:
```
docker-compose -f docker-compose-vertexai.yaml down --volumes
```
To confirm all Docker containers have been shut down, check that the following list is empty:
```
docker ps
```
To confirm all Docker volumes have been removed, check that the following list is empty:
```
docker volume ls
```

6
scripts/.ipynb_checkpoints/kg-extract-definitions-checkpoint Executable file
View file

@ -0,0 +1,6 @@
#!/usr/bin/env python3
from trustgraph.kg.extract_definitions import run
run()

6
scripts/.ipynb_checkpoints/llm-azure-text-checkpoint Executable file
View file

@ -0,0 +1,6 @@
#!/usr/bin/env python3
from trustgraph.llm.azure_text import run
run()

6
scripts/text-completion-cohere Executable file
View file

@ -0,0 +1,6 @@
#!/usr/bin/env python3
from trustgraph.model.text_completion.cohere import run
run()

BIN
sources/Challenger-Report-Vol1.pdf Normal file
View file

Binary file not shown.

17301
sources/solarwinds-dismissal-opinion.pdf Normal file
View file

File diff suppressed because it is too large Load diff

71
trustgraph/.ipynb_checkpoints/llm_client-checkpoint.py Normal file
View file

@ -0,0 +1,71 @@
#!/usr/bin/env python3
import pulsar
import _pulsar
from pulsar.schema import JsonSchema
from trustgraph.schema import TextCompletionRequest, TextCompletionResponse
import hashlib
import uuid
# Ugly
ERROR=_pulsar.LoggerLevel.Error
WARN=_pulsar.LoggerLevel.Warn
INFO=_pulsar.LoggerLevel.Info
DEBUG=_pulsar.LoggerLevel.Debug
class LlmClient:
def __init__(
self, log_level=ERROR, client_id=None,
pulsar_host="pulsar://pulsar:6650",
):
if client_id == None:
client_id = str(uuid.uuid4())
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level),
)
self.producer = self.client.create_producer(
topic='llm-complete-text',
schema=JsonSchema(TextCompletionRequest),
chunking_enabled=True,
)
self.consumer = self.client.subscribe(
'llm-complete-text-response', client_id,
schema=JsonSchema(TextCompletionResponse),
)
def request(self, prompt, timeout=500):
id = str(uuid.uuid4())
r = TextCompletionRequest(
prompt=prompt
)
self.producer.send(r, properties={ "id": id })
while True:
msg = self.consumer.receive(timeout_millis=timeout * 1000)
mid = msg.properties()["id"]
if mid == id:
resp = msg.value().response
self.consumer.acknowledge(msg)
return resp
# Ignore messages with wrong ID
self.consumer.acknowledge(msg)
def __del__(self):
self.producer.close()
self.consumer.close()
self.client.close()

138
trustgraph/.ipynb_checkpoints/prompts-checkpoint.py Normal file
View file

@ -0,0 +1,138 @@
def turtle_extract(text):
prompt = f"""<instructions>
Study the following text and extract knowledge as
information in Turtle RDF format.
When declaring any new URIs, use <https://trustgraph.ai/e#> prefix,
and declare appropriate namespace tags.
</instructions>
<text>
{text}
</text>
<requirements>
Do not use placeholders for information you do not know.
You will respond only with raw Turtle RDF data. Do not provide
explanations. Do not use special characters in the abstract text. The
abstract must be written as plain text. Do not add markdown formatting.
</requirements>"""
return prompt
def scholar(text):
# Build the prompt for Article style extraction
jsonexample = """{
"title": "Article title here",
"abstract": "Abstract text here",
"keywords": ["keyword1", "keyword2", "keyword3"],
"people": ["person1", "person2", "person3"]
}"""
promptscholar = f"""Your task is to read the provided text and write a scholarly abstract to fully explain all of the concepts described in the provided text. The abstract must include all conceptual details.
<text>
{text}
</text>
<instructions>
- Structure: For the provided text, write a title, abstract, keywords,
and people for the concepts found in the provided text. Ignore
document formatting in the provided text such as table of contents,
headers, footers, section metadata, and URLs.
- Focus on Concepts The abstract must focus on concepts found in the
provided text. The abstract must be factually accurate. Do not
write any concepts not found in the provided text. Do not
speculate. Do not omit any conceptual details.
- Completeness: The abstract must capture all topics the reader will
need to understand the concepts found in the provided text. Describe
all terms, definitions, entities, people, events, concepts,
conceptual relationships, and any other topics necessary for the
reader to understand the concepts of the provided text.
- Format: Respond in the form of a valid JSON object.
</instructions>
<example>
{jsonexample}
</example>
<requirements>
You will respond only with the JSON object. Do not provide
explanations. Do not use special characters in the abstract text. The
abstract must be written as plain text.
</requirements>"""
return promptscholar
def to_json_ld(text):
prompt = f"""<instructions>
Study the following text and output any facts you discover in
well-structured JSON-LD format.
Use any schema you understand from schema.org to describe the facts.
</instructions>
<text>
{text}
</text>
<requirements>
You will respond only with raw JSON-LD data in JSON format. Do not provide
explanations. Do not use special characters in the abstract text. The
abstract must be written as plain text. Do not add markdown formatting
or headers or prefixes. Do not use information which is not present in
the input text.
</requirements>"""
return prompt
def to_relationships(text):
prompt = f"""<instructions>
Study the following text and derive entity relationships. For each
relationship, derive the subject, predicate and object of the relationship.
Output relationships in JSON format as an arary of objects with fields:
- subject: the subject of the relationship
- predicate: the predicate
- object: the object of the relationship
- object-entity: false if the object is a simple data type: name, value or date. true if it is an entity.
</instructions>
<text>
{text}
</text>
<requirements>
You will respond only with raw JSON format data. Do not provide
explanations. Do not use special characters in the abstract text. The
abstract must be written as plain text. Do not add markdown formatting
or headers or prefixes.
</requirements>"""
return prompt
def to_definitions(text):
prompt = f"""<instructions>
Study the following text and derive definitions for any discovered entities.
Do not provide definitions for entities whose definitions are incomplete
or unknown.
Output relationships in JSON format as an arary of objects with fields:
- entity: the name of the entity
- definition: English text which defines the entity
</instructions>
<text>
{text}
</text>
<requirements>
You will respond only with raw JSON format data. Do not provide
explanations. Do not use special characters in the abstract text. The
abstract will be written as plain text. Do not add markdown formatting
or headers or prefixes. Do not include null or unknown definitions.
</requirements>"""
return prompt

6
trustgraph/.ipynb_checkpoints/rdf-checkpoint.py Normal file
View file

@ -0,0 +1,6 @@
RDF_LABEL = "http://www.w3.org/2000/01/rdf-schema#label"
DEFINITION = "http://www.w3.org/2004/02/skos/core#definition"
TRUSTGRAPH_ENTITIES = "http://trustgraph.ai/e/"

3
trustgraph/kg/extract_definitions/.ipynb_checkpoints/__init__-checkpoint.py Normal file
View file

@ -0,0 +1,3 @@
from . extract import *

7
trustgraph/kg/extract_definitions/.ipynb_checkpoints/__main__-checkpoint.py Executable file
View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . extract import run
if __name__ == '__main__':
run()

197
trustgraph/kg/extract_definitions/.ipynb_checkpoints/extract-checkpoint.py Executable file
View file

@ -0,0 +1,197 @@
"""
Simple decoder, accepts vector+text chunks input, applies entity analysis to
get entity definitions which are output as graph edges.
"""
import pulsar
from pulsar.schema import JsonSchema
from langchain_community.document_loaders import PyPDFLoader
import tempfile
import base64
import os
import argparse
import rdflib
import json
import urllib.parse
import time
from ... schema import VectorsChunk, Triple, Source, Value
from ... log_level import LogLevel
from ... llm_client import LlmClient
from ... prompts import to_definitions
from ... rdf import TRUSTGRAPH_ENTITIES, DEFINITION
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
default_input_queue = 'vectors-chunk-load'
default_output_queue = 'graph-load'
default_subscriber = 'kg-extract-definitions'
class Processor:
def __init__(
self,
pulsar_host=default_pulsar_host,
input_queue=default_input_queue,
output_queue=default_output_queue,
subscriber=default_subscriber,
log_level=LogLevel.INFO,
):
self.client = None
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.consumer = self.client.subscribe(
input_queue, subscriber,
schema=JsonSchema(VectorsChunk),
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(Triple),
)
self.llm = LlmClient(pulsar_host=pulsar_host)
def to_uri(self, text):
part = text.replace(" ", "-").lower().encode("utf-8")
quoted = urllib.parse.quote(part)
uri = TRUSTGRAPH_ENTITIES + quoted
return uri
def get_definitions(self, chunk):
prompt = to_definitions(chunk)
resp = self.llm.request(prompt)
defs = json.loads(resp)
return defs
def emit_edge(self, s, p, o):
t = Triple(s=s, p=p, o=o)
self.producer.send(t)
def run(self):
while True:
msg = self.consumer.receive()
try:
v = msg.value()
print(f"Indexing {v.source.id}...", flush=True)
chunk = v.chunk.decode("utf-8")
g = rdflib.Graph()
try:
defs = self.get_definitions(chunk)
print(json.dumps(defs, indent=4), flush=True)
for defn in defs:
s = defn["entity"]
s_uri = self.to_uri(s)
o = defn["definition"]
s_value = Value(value=str(s_uri), is_uri=True)
o_value = Value(value=str(o), is_uri=False)
self.emit_edge(s_value, DEFINITION_VALUE, o_value)
except Exception as e:
print("Exception: ", e, flush=True)
print("Done.", flush=True)
# Acknowledge successful processing of the message
self.consumer.acknowledge(msg)
except Exception as e:
print("Exception: ", e, flush=True)
# Message failed to be processed
self.consumer.negative_acknowledge(msg)
def __del__(self):
if self.client:
self.client.close()
def run():
parser = argparse.ArgumentParser(
prog='pdf-decoder',
description=__doc__,
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
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})'
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.INFO,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
args = parser.parse_args()
while True:
try:
p = Processor(
pulsar_host=args.pulsar_host,
input_queue=args.input_queue,
output_queue=args.output_queue,
subscriber=args.subscriber,
log_level=args.log_level,
)
p.run()
except Exception as e:
print("Exception:", e, flush=True)
print("Will retry...", flush=True)
time.sleep(10)

3
trustgraph/kg/extract_relationships/.ipynb_checkpoints/__init__-checkpoint.py Normal file
View file

@ -0,0 +1,3 @@
from . extract import *

7
trustgraph/kg/extract_relationships/.ipynb_checkpoints/__main__-checkpoint.py Executable file
View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . extract import run
if __name__ == '__main__':
run()

253
trustgraph/kg/extract_relationships/.ipynb_checkpoints/extract-checkpoint.py Executable file
View file

@ -0,0 +1,253 @@
"""
Simple decoder, accepts vector+text chunks input, applies entity
relationship analysis to get entity relationship edges which are output as
graph edges.
"""
import pulsar
from pulsar.schema import JsonSchema
from langchain_community.document_loaders import PyPDFLoader
import tempfile
import base64
import os
import argparse
import rdflib
import json
import urllib.parse
import time
from ... schema import VectorsChunk, Triple, VectorsAssociation, Source, Value
from ... log_level import LogLevel
from ... llm_client import LlmClient
from ... prompts import to_relationships
from ... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
default_input_queue = 'vectors-chunk-load'
default_output_queue = 'graph-load'
default_subscriber = 'kg-extract-relationships'
default_vector_queue='vectors-load'
class Processor:
def __init__(
self,
pulsar_host=default_pulsar_host,
input_queue=default_input_queue,
vector_queue=default_vector_queue,
output_queue=default_output_queue,
subscriber=default_subscriber,
log_level=LogLevel.INFO,
):
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.consumer = self.client.subscribe(
input_queue, subscriber,
schema=JsonSchema(VectorsChunk),
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(Triple),
)
self.vec_prod = self.client.create_producer(
topic=vector_queue,
schema=JsonSchema(VectorsAssociation),
)
self.llm = LlmClient(pulsar_host=pulsar_host)
def to_uri(self, text):
part = text.replace(" ", "-").lower().encode("utf-8")
quoted = urllib.parse.quote(part)
uri = TRUSTGRAPH_ENTITIES + quoted
return uri
def get_relationships(self, chunk):
prompt = to_relationships(chunk)
resp = self.llm.request(prompt)
rels = json.loads(resp)
return rels
def emit_edge(self, s, p, o):
t = Triple(s=s, p=p, o=o)
self.producer.send(t)
def emit_vec(self, ent, vec):
r = VectorsAssociation(entity=ent, vectors=vec)
self.vec_prod.send(r)
def run(self):
while True:
msg = self.consumer.receive()
try:
v = msg.value()
print(f"Indexing {v.source.id}...", flush=True)
chunk = v.chunk.decode("utf-8")
g = rdflib.Graph()
try:
rels = self.get_relationships(chunk)
print(json.dumps(rels, indent=4), flush=True)
for rel in rels:
s = rel["subject"]
p = rel["predicate"]
o = rel["object"]
s_uri = self.to_uri(s)
s_value = Value(value=str(s_uri), is_uri=True)
p_uri = self.to_uri(p)
p_value = Value(value=str(p_uri), is_uri=True)
if rel["object-entity"]:
o_uri = self.to_uri(o)
o_value = Value(value=str(o_uri), is_uri=True)
else:
o_value = Value(value=str(o), is_uri=False)
self.emit_edge(
s_value,
p_value,
o_value
)
# Label for s
self.emit_edge(
s_value,
RDF_LABEL_VALUE,
Value(value=str(s), is_uri=False)
)
# Label for p
self.emit_edge(
p_value,
RDF_LABEL_VALUE,
Value(value=str(p), is_uri=False)
)
if rel["object-entity"]:
# Label for o
self.emit_edge(
o_value,
RDF_LABEL_VALUE,
Value(value=str(o), is_uri=False)
)
self.emit_vec(s_value, v.vectors)
self.emit_vec(p_value, v.vectors)
if rel["object-entity"]:
self.emit_vec(o_value, v.vectors)
except Exception as e:
print("Exception: ", e, flush=True)
print("Done.", flush=True)
# Acknowledge successful processing of the message
self.consumer.acknowledge(msg)
except Exception as e:
print("Exception: ", e, flush=True)
# Message failed to be processed
self.consumer.negative_acknowledge(msg)
def __del__(self):
self.client.close()
def run():
parser = argparse.ArgumentParser(
prog='kg-extract-relationships',
description=__doc__,
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
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})'
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.INFO,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'-c', '--vector-queue',
default=default_vector_queue,
help=f'Vector output queue (default: {default_vector_queue})'
)
args = parser.parse_args()
while True:
try:
p = Processor(
pulsar_host=args.pulsar_host,
input_queue=args.input_queue,
output_queue=args.output_queue,
vector_queue=args.vector_queue,
subscriber=args.subscriber,
log_level=args.log_level,
)
p.run()
except Exception as e:
print("Exception:", e, flush=True)
print("Will retry...", flush=True)
time.sleep(10)

127
trustgraph/llm/azure_text/.ipynb_checkpoints/llm-checkpoint.py Executable file
View file

@ -0,0 +1,127 @@
"""
Simple LLM service, performs text prompt completion using the Azure
serverless endpoint service. Input is prompt, output is response.
"""
import requests
import json
from ... schema import TextCompletionRequest, TextCompletionResponse
from ... log_level import LogLevel
from ... base import ConsumerProducer
default_input_queue = 'llm-complete-text'
default_output_queue = 'llm-complete-text-response'
default_subscriber = 'llm-azure-text'
class Processor(ConsumerProducer):
def __init__(
self,
pulsar_host=None,
input_queue=default_input_queue,
output_queue=default_output_queue,
subscriber=default_subscriber,
log_level=LogLevel.INFO,
endpoint=None,
token=None,
):
super(Processor, self).__init__(
pulsar_host=pulsar_host,
log_level=log_level,
input_queue=input_queue,
output_queue=output_queue,
subscriber=subscriber,
input_schema=TextCompletionRequest,
output_schema=TextCompletionResponse,
)
self.endpoint = endpoint
self.token = token
def build_prompt(self, system, content):
data = {
"messages": [
{
"role": "system", "content": system
},
{
"role": "user", "content": content
}
],
"max_tokens": 4192,
"temperature": 0.2,
"top_p": 1
}
body = json.dumps(data)
return body
def call_llm(self, body):
url = self.endpoint
# Replace this with the primary/secondary key, AMLToken, or
# Microsoft Entra ID token for the endpoint
api_key = self.token
headers = {
'Content-Type': 'application/json',
'Authorization': f'Bearer {api_key}'
}
resp = requests.post(url, data=body, headers=headers)
result = resp.json()
message_content = result['choices'][0]['message']['content']
return message_content
def handle(self, msg):
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling prompt {id}...", flush=True)
prompt = self.build_prompt(
"You are a helpful chatbot",
v.prompt
)
response = self.call_llm(prompt)
print("Send response...", flush=True)
r = TextCompletionResponse(response=response)
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(
'-e', '--endpoint',
help=f'LLM model endpoint'
)
parser.add_argument(
'-k', '--token',
help=f'LLM model token'
)
def run():
Processor.start("llm-azure-text", __doc__)

192
trustgraph/llm/claude_text/.ipynb_checkpoints/llm-checkpoint.py Executable file
View file

@ -0,0 +1,192 @@
"""
Simple LLM service, performs text prompt completion using Claude.
Input is prompt, output is response.
"""
import pulsar
from pulsar.schema import JsonSchema
import tempfile
import base64
import os
import argparse
import anthropic
import time
from ... schema import TextCompletionRequest, TextCompletionResponse
from ... log_level import LogLevel
default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
default_input_queue = 'llm-complete-text'
default_output_queue = 'llm-complete-text-response'
default_subscriber = 'llm-claude-text'
default_model = 'claude-3-5-sonnet-20240620'
class Processor:
def __init__(
self,
pulsar_host=default_pulsar_host,
input_queue=default_input_queue,
output_queue=default_output_queue,
subscriber=default_subscriber,
log_level=LogLevel.INFO,
model=default_model,
api_key,
):
self.client = None
self.client = pulsar.Client(
pulsar_host,
logger=pulsar.ConsoleLogger(log_level.to_pulsar())
)
self.consumer = self.client.subscribe(
input_queue, subscriber,
schema=JsonSchema(TextCompletionRequest),
)
self.producer = self.client.create_producer(
topic=output_queue,
schema=JsonSchema(TextCompletionResponse),
)
self.model = model
self.claude = anthropic.Anthropic(api_key=api_key)
print("Initialised", flush=True)
def run(self):
while True:
msg = self.consumer.receive()
try:
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling prompt {id}...", flush=True)
prompt = v.prompt
response = message = self.claude.messages.create(
model=self.model,
max_tokens=1000,
temperature=0.1,
system = "You are a helpful chatbot.",
messages=[
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt
}
]
}
]
)
resp = response.content[0].text
print(resp, flush=True)
print("Send response...", flush=True)
r = TextCompletionResponse(response=resp)
self.producer.send(r, properties={"id": id})
print("Done.", flush=True)
# Acknowledge successful processing of the message
self.consumer.acknowledge(msg)
except Exception as e:
print("Exception:", e, flush=True)
# Message failed to be processed
self.consumer.negative_acknowledge(msg)
def __del__(self):
self.client.close()
def run():
parser = argparse.ArgumentParser(
prog='llm-ollama-text',
description=__doc__,
)
parser.add_argument(
'-p', '--pulsar-host',
default=default_pulsar_host,
help=f'Pulsar host (default: {default_pulsar_host})',
)
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})'
)
parser.add_argument(
'-o', '--output-queue',
default=default_output_queue,
help=f'Output queue (default: {default_output_queue})'
)
parser.add_argument(
'-l', '--log-level',
type=LogLevel,
default=LogLevel.INFO,
choices=list(LogLevel),
help=f'Output queue (default: info)'
)
parser.add_argument(
'-m', '--model',
default="claude-3-5-sonnet-20240620",
help=f'LLM model (default: claude-3-5-sonnet-20240620)'
)
parser.add_argument(
'-k', '--api-key',
help=f'Claude API key'
)
args = parser.parse_args()
while True:
try:
p = Processor(
pulsar_host=args.pulsar_host,
input_queue=args.input_queue,
output_queue=args.output_queue,
subscriber=args.subscriber,
log_level=args.log_level,
model=args.model,
api_key=args.api_key,
)
p.run()
except Exception as e:
print("Exception:", e, flush=True)
print("Will retry...", flush=True)
time.sleep(10)

3
trustgraph/model/text_completion/cohere/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from . llm import *

7
trustgraph/model/text_completion/cohere/__main__.py Executable file
View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . llm import run
if __name__ == '__main__':
run()

105
trustgraph/model/text_completion/cohere/llm.py Executable file
View file

@ -0,0 +1,105 @@
"""
Simple LLM service, performs text prompt completion using Cohere.
Input is prompt, output is response.
"""
import cohere
from .... schema import TextCompletionRequest, TextCompletionResponse
from .... schema import text_completion_request_queue
from .... schema import text_completion_response_queue
from .... log_level import LogLevel
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue
default_subscriber = module
default_model = 'c4ai-aya-23-8b'
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)
api_key = params.get("api_key")
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": TextCompletionRequest,
"output_schema": TextCompletionResponse,
"model": model,
}
)
self.model = model
self.cohere = cohere.Client(api_key=api_key)
print("Initialised", flush=True)
def handle(self, msg):
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling prompt {id}...", flush=True)
prompt = v.prompt
stream = self.cohere.chat_stream(
model=self.model,
#model='c4ai-aya-23-8b',
message=prompt,
temperature=0.0,
chat_history=[],
prompt_truncation='auto',
connectors=[]
)
for event in stream:
if event.event_type == "text-generation":
resp = event.text
print(resp, flush=True)
print("Send response...", flush=True)
r = TextCompletionResponse(response=resp)
self.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="c4ai-aya-23-8b",
help=f'Cohere model (default: c4ai-aya-23-8b)'
)
parser.add_argument(
'-k', '--api-key',
help=f'Cohere API key'
)
def run():
Processor.start(module, __doc__)