Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Amazon Bedrock support ⛰️ #49

Merged
merged 2 commits into from
Oct 19, 2024
Merged

Conversation

JGalego
Copy link
Contributor

@JGalego JGalego commented Oct 18, 2024

This PR adds support for Amazon Bedrock LLMs via Converse API and Embedding models (Titan + Cohere) with InvokeModel. Implementation is based on async calls via aioboto3 cf. terricain/aioboto3#341 for additional info.

Example:

"""
LightRAG meets Amazon Bedrock ⛰️
"""

import os
import logging

from lightrag import LightRAG, QueryParam
from lightrag.llm import bedrock_complete, bedrock_embedding
from lightrag.utils import EmbeddingFunc

logging.getLogger("aiobotocore").setLevel(logging.WARNING)

WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)

rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=bedrock_complete,
    llm_model_name="Anthropic Claude 3 Haiku // Amazon Bedrock",
    embedding_func=EmbeddingFunc(
        embedding_dim=1024,
        max_token_size=8192,
        func=bedrock_embedding
    )
)

with open("./book.txt", 'r', encoding='utf-8') as f:
    rag.insert(f.read())

for mode in ["naive", "local", "global", "hybrid"]:
    print("\n+-" + "-" * len(mode) + "-+")
    print(f"| {mode.capitalize()} |")
    print("+-" + "-" * len(mode) + "-+\n")
    print(
        rag.query(
            "What are the top themes in this story?",
            param=QueryParam(mode=mode)
        )
    )

Output:

INFO:lightrag:Logger initialized for working directory: ./dickens
DEBUG:lightrag:LightRAG init with param:
  working_dir = ./dickens,
  chunk_token_size = 1200,
  chunk_overlap_token_size = 100,
  tiktoken_model_name = gpt-4o-mini,
  entity_extract_max_gleaning = 1,
  entity_summary_to_max_tokens = 500,
  node_embedding_algorithm = node2vec,
  node2vec_params = {'dimensions': 1536, 'num_walks': 10, 'walk_length': 40, 'window_size': 2, 'iterations': 3, 'random_seed': 3},
  embedding_func = {'embedding_dim': 1024, 'max_token_size': 8192, 'func': <function <lambda> at 0x7fcca60096c0>},
  embedding_batch_num = 32,
  embedding_func_max_async = 16,
  llm_model_func = <function bedrock_complete at 0x7fcca4ad7100>,
  llm_model_name = Anthropic Claude 3 Haiku // Amazon Bedrock,
  llm_model_max_token_size = 32768,
  llm_model_max_async = 16,
  key_string_value_json_storage_cls = <class 'lightrag.storage.JsonKVStorage'>,
  vector_db_storage_cls = <class 'lightrag.storage.NanoVectorDBStorage'>,
  vector_db_storage_cls_kwargs = {},
  graph_storage_cls = <class 'lightrag.storage.NetworkXStorage'>,
  enable_llm_cache = True,
  addon_params = {},
  convert_response_to_json_func = <function convert_response_to_json at 0x7fcca4ad45e0>

INFO:lightrag:Load KV full_docs with 0 data
INFO:lightrag:Load KV text_chunks with 0 data
INFO:lightrag:Load KV llm_response_cache with 0 data
INFO:nano-vectordb:Init {'embedding_dim': 1024, 'metric': 'cosine', 'storage_file': './dickens/vdb_entities.json'} 0 data
INFO:nano-vectordb:Init {'embedding_dim': 1024, 'metric': 'cosine', 'storage_file': './dickens/vdb_relationships.json'} 0 data
INFO:nano-vectordb:Init {'embedding_dim': 1024, 'metric': 'cosine', 'storage_file': './dickens/vdb_chunks.json'} 0 data
INFO:lightrag:Creating a new event loop in a sub-thread.
INFO:lightrag:[New Docs] inserting 1 docs
INFO:lightrag:[New Chunks] inserting 42 chunks
INFO:lightrag:Inserting 42 vectors to chunks
INFO:lightrag:[Entity Extraction]...
⠹ Processed 42 chunks, 415 entities(duplicated), 261 relations(duplicated)
DEBUG:lightrag:Trigger summary: "SCROOGE"
INFO:lightrag:Inserting 264 vectors to entities
INFO:lightrag:Inserting 215 vectors to relationships
INFO:lightrag:Writing graph with 281 nodes, 215 edges

+-------+
| Naive |
+-------+

INFO:lightrag:Creating a new event loop in a sub-thread.
Sorry, I'm not able to provide an answer to that question.

+-------+
| Local |
+-------+

INFO:lightrag:Creating a new event loop in a sub-thread.
INFO:lightrag:Local query uses 1 entites, 1 relations, 1 text units
Based on the information provided in the data tables, the top themes in this story appear to be:

## Symbolism of Wealth and Material Possessions
The data indicates that the keys that are part of the chain binding Marley's ghost represent his focus on material wealth and possessions during his life. This symbolism suggests a theme about the dangers of being overly consumed by material concerns and neglecting more spiritual or ethical aspects of life.

## Guilt, Regret, and the Afterlife
Marley's ghost appears to Scrooge to warn him and impart a moral lesson. The fact that Marley is doomed to wander the earth bound by a heavy chain made of "cash-boxes, keys, padlocks, ledgers, deeds, and heavy purses" indicates a theme of guilt and regret over one's actions in life and the consequences in the afterlife.

## The Nature of Reality and Skepticism
Scrooge's initial skepticism and disbelief towards Marley's ghost, questioning whether it is real or just "a disorder of the stomach", suggests a theme around the nature of reality, the limits of our senses, and a skeptical worldview. This ties into the ghost's challenge to Scrooge to consider the evidence of his own senses.

Without more context from the full story, I cannot say definitively that these are the top themes. But based on the information provided, these appear to be the key thematic elements present. Let me know if you need any clarification or have additional questions!

+--------+
| Global |
+--------+

INFO:lightrag:Creating a new event loop in a sub-thread.
INFO:lightrag:Global query uses 9 entites, 6 relations, 3 text units
Based on the information provided in the data tables, the top themes in this story appear to be:

## Symbolism of Wealth and Material Possessions
The data indicates that the keys that are part of the chain binding Marley's ghost represent his focus on material wealth and possessions during his life. This symbolism suggests a theme about the dangers of being overly consumed by material concerns and neglecting more spiritual or ethical aspects of life.

## Guilt, Regret, and the Afterlife
Marley's ghost appears to Scrooge to warn him and impart a moral lesson. The fact that Marley is doomed to wander the earth bound by a heavy chain made of "cash-boxes, keys, padlocks, ledgers, deeds, and heavy purses" indicates a theme of guilt and regret over one's actions in life and the consequences in the afterlife.

## The Nature of Reality and Skepticism
Scrooge's initial skepticism and disbelief towards Marley's ghost, questioning whether it is real or just "a disorder of the stomach", suggests a theme around the nature of reality, the limits of our senses, and a skeptical worldview. This ties into the ghost's challenge to Scrooge to consider the evidence of his own senses.

Without more context from the full story, I cannot say definitively that these are the top themes. But based on the information provided, these appear to be the key thematic elements present. Let me know if you need any clarification or have additional questions!

+--------+
| Hybrid |
+--------+

INFO:lightrag:Creating a new event loop in a sub-thread.
INFO:lightrag:Local query uses 1 entites, 1 relations, 1 text units
INFO:lightrag:Global query uses 9 entites, 6 relations, 3 text units
Based on the information provided in the data tables, the top themes in this story appear to be:

## Symbolism of Wealth and Material Possessions
The data indicates that the keys that are part of the chain binding Marley's ghost represent his focus on material wealth and possessions during his life. This symbolism suggests a theme about the dangers of being overly consumed by material concerns and neglecting more spiritual or ethical aspects of life.

## Guilt, Regret, and the Afterlife
Marley's ghost appears to Scrooge to warn him and impart a moral lesson. The fact that Marley is doomed to wander the earth bound by a heavy chain made of "cash-boxes, keys, padlocks, ledgers, deeds, and heavy purses" indicates a theme of guilt and regret over one's actions in life and the consequences in the afterlife.

## The Nature of Reality and Skepticism
Scrooge's initial skepticism and disbelief towards Marley's ghost, questioning whether it is real or just "a disorder of the stomach", suggests a theme around the nature of reality, the limits of our senses, and a skeptical worldview. This ties into the ghost's challenge to Scrooge to consider the evidence of his own senses.

Without more context from the full story, I cannot say definitively that these are the top themes. But based on the information provided, these appear to be the key thematic elements present. Let me know if you need any clarification or have additional questions!

133167be-cdd8-4f26-9678-758e1546c2e4_text

@JGalego JGalego marked this pull request as draft October 18, 2024 13:26
@JGalego JGalego marked this pull request as ready for review October 18, 2024 15:50
@LarFii
Copy link
Collaborator

LarFii commented Oct 19, 2024

Thanks for your contribution! Best regards!

@LarFii LarFii merged commit b854ab4 into HKUDS:main Oct 19, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants