Merge pull request #39 from TogetherCrew/feat/llama-index-updates

feat: Added update based on new lib changes!
TogetherCrew · Mar 13, 2024 · 5265001 · 5265001
2 parents 3df30b6 + 651c182
commit 5265001
Show file tree

Hide file tree

Showing 19 changed files with 56 additions and 53 deletions.
diff --git a/bot/retrievers/forum_summary_retriever.py b/bot/retrievers/forum_summary_retriever.py
@@ -1,6 +1,6 @@
 from bot.retrievers.summary_retriever_base import BaseSummarySearch
-from llama_index.embeddings import BaseEmbedding
-from llama_index.schema import NodeWithScore
+from llama_index.core.embeddings import BaseEmbedding
+from llama_index.core.schema import NodeWithScore
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 
 

diff --git a/bot/retrievers/retrieve_similar_nodes.py b/bot/retrievers/retrieve_similar_nodes.py
@@ -1,10 +1,11 @@
 from datetime import datetime, timedelta
 
 from dateutil import parser
-from llama_index.embeddings import BaseEmbedding
-from llama_index.schema import NodeWithScore
-from llama_index.vector_stores import PGVectorStore, VectorStoreQueryResult
-from llama_index.vector_stores.postgres import DBEmbeddingRow
+from llama_index.core.embeddings import BaseEmbedding
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.vector_stores.types import VectorStoreQueryResult
+from llama_index.vector_stores.postgres import PGVectorStore
+from llama_index.vector_stores.postgres.base import DBEmbeddingRow
 from sqlalchemy import Date, and_, cast, null, or_, select, text
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 

diff --git a/bot/retrievers/summary_retriever_base.py b/bot/retrievers/summary_retriever_base.py
@@ -1,7 +1,6 @@
-from llama_index import VectorStoreIndex
-from llama_index.embeddings import BaseEmbedding
-from llama_index.indices.query.schema import QueryBundle
-from llama_index.schema import NodeWithScore
+from llama_index.core import VectorStoreIndex
+from llama_index.core.embeddings import BaseEmbedding
+from llama_index.core.schema import NodeWithScore, QueryBundle
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 from tc_hivemind_backend.pg_vector_access import PGVectorAccess
 

diff --git a/discord_query.py b/discord_query.py
@@ -1,5 +1,5 @@
-from llama_index import QueryBundle
-from llama_index.schema import NodeWithScore
+from llama_index.core import QueryBundle
+from llama_index.core.schema import NodeWithScore
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 from utils.query_engine.prepare_discord_query_engine import (
     prepare_discord_engine_auto_filter,

diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
-python3 -m coverage run --omit=tests/* -m pytest .
+python3 -m coverage run --omit=tests/* -m pytest . && echo "Tests Passed" || exit 1
 python3 -m coverage lcov -o coverage/lcov.info
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,4 @@
 numpy
-llama-index==0.9.48
 pymongo
 python-dotenv
 pgvector
@@ -16,7 +15,9 @@ neo4j>=5.14.1, <6.0.0
 coverage>=7.3.3, <8.0.0
 pytest>=7.4.3, <8.0.0
 python-dotenv==1.0.0
-tc-hivemind-backend==1.1.2
+tc-hivemind-backend==1.1.5
+llama-index-question-gen-guidance==0.1.2
+llama-index-vector-stores-postgres==0.1.2
 celery>=5.3.6, <6.0.0
 guidance
 tc-messageBroker==1.6.6

diff --git a/subquery.py b/subquery.py
@@ -1,10 +1,11 @@
 from guidance.models import OpenAIChat
-from llama_index import QueryBundle, ServiceContext
-from llama_index.core.base_query_engine import BaseQueryEngine
-from llama_index.query_engine import SubQuestionQueryEngine
-from llama_index.question_gen.guidance_generator import GuidanceQuestionGenerator
-from llama_index.schema import NodeWithScore
-from llama_index.tools import QueryEngineTool, ToolMetadata
+from llama_index.core import QueryBundle, Settings
+from llama_index.core.base.base_query_engine import BaseQueryEngine
+from llama_index.core.query_engine import SubQuestionQueryEngine
+from llama_index.core.schema import NodeWithScore
+from llama_index.core.tools import QueryEngineTool, ToolMetadata
+from llama_index.llms.openai import OpenAI
+from llama_index.question_gen.guidance import GuidanceQuestionGenerator
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 from utils.query_engine import prepare_discord_engine_auto_filter
 
@@ -97,17 +98,19 @@ def query_multiple_source(
     if github:
         raise NotImplementedError
 
+    embed_model = CohereEmbedding()
+    llm = OpenAI("gpt-3.5-turbo")
+    Settings.embed_model = embed_model
+    Settings.llm = llm
+
     question_gen = GuidanceQuestionGenerator.from_defaults(
         guidance_llm=OpenAIChat("gpt-4"),
         verbose=False,
     )
-    embed_model = CohereEmbedding()
-    service_context = ServiceContext.from_defaults(embed_model=embed_model)
     s_engine = SubQuestionQueryEngine.from_defaults(
         question_gen=question_gen,
         query_engine_tools=query_engine_tools,
         use_async=False,
-        service_context=service_context,
     )
     query_embedding = embed_model.get_text_embedding(text=query)
     response = s_engine.query(QueryBundle(query_str=query, embedding=query_embedding))

diff --git a/tests/integration/test_retrieve_similar_nodes.py b/tests/integration/test_retrieve_similar_nodes.py
@@ -2,7 +2,7 @@
 from unittest.mock import MagicMock
 
 from bot.retrievers.retrieve_similar_nodes import RetrieveSimilarNodes
-from llama_index.schema import TextNode
+from llama_index.core.schema import TextNode
 
 
 class TestRetrieveSimilarNodes(TestCase):

diff --git a/tests/unit/test_discord_summary_retriever.py b/tests/unit/test_discord_summary_retriever.py
@@ -5,7 +5,7 @@
 
 from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever
 from dateutil import parser
-from llama_index import Document, MockEmbedding, ServiceContext, VectorStoreIndex
+from llama_index.core import Document, MockEmbedding, Settings, VectorStoreIndex
 
 
 class TestDiscordSummaryRetriever(TestCase):
@@ -31,13 +31,10 @@ def test_initialize_class(self):
 
         mock_embedding_model = partial(MockEmbedding, embed_dim=1024)
 
-        service_context = ServiceContext.from_defaults(
-            llm=None, embed_model=mock_embedding_model()
-        )
+        Settings.llm = None
+        Settings.embed_model = mock_embedding_model()
         ForumBasedSummaryRetriever._setup_index.return_value = (
-            VectorStoreIndex.from_documents(
-                documents=[doc], service_context=service_context
-            )
+            VectorStoreIndex.from_documents(documents=[doc])
         )
 
         base_summary_search = ForumBasedSummaryRetriever(

diff --git a/tests/unit/test_level_based_platform_query_engine.py b/tests/unit/test_level_based_platform_query_engine.py
@@ -4,7 +4,7 @@
 
 from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever
 from bot.retrievers.retrieve_similar_nodes import RetrieveSimilarNodes
-from llama_index.schema import NodeWithScore, TextNode
+from llama_index.core.schema import NodeWithScore, TextNode
 from sqlalchemy.exc import OperationalError
 from utils.query_engine.level_based_platform_query_engine import (
     LevelBasedPlatformQueryEngine,

diff --git a/tests/unit/test_level_based_platform_query_engine_prepare_context.py b/tests/unit/test_level_based_platform_query_engine_prepare_context.py
@@ -3,7 +3,7 @@
 from unittest.mock import patch
 
 from bot.retrievers.retrieve_similar_nodes import RetrieveSimilarNodes
-from llama_index.schema import NodeWithScore, TextNode
+from llama_index.core.schema import NodeWithScore, TextNode
 from utils.query_engine.level_based_platform_query_engine import (
     LevelBasedPlatformQueryEngine,
 )

diff --git a/tests/unit/test_level_based_platform_util.py b/tests/unit/test_level_based_platform_util.py
@@ -1,6 +1,6 @@
 import unittest
 
-from llama_index.schema import NodeWithScore, TextNode
+from llama_index.core.schema import NodeWithScore, TextNode
 from utils.query_engine.level_based_platforms_util import LevelBasedPlatformUtils
 
 

diff --git a/tests/unit/test_prepare_discord_query_engine.py b/tests/unit/test_prepare_discord_query_engine.py
@@ -1,7 +1,7 @@
 import os
 import unittest
 
-from llama_index.core.base_query_engine import BaseQueryEngine
+from llama_index.core.base.base_query_engine import BaseQueryEngine
 from utils.query_engine.prepare_discord_query_engine import prepare_discord_engine
 
 

diff --git a/tests/unit/test_prepare_discourse_query_engine.py b/tests/unit/test_prepare_discourse_query_engine.py
@@ -1,7 +1,7 @@
 import os
 import unittest
 
-from llama_index.core.base_query_engine import BaseQueryEngine
+from llama_index.core.base.base_query_engine import BaseQueryEngine
 from utils.query_engine.prepare_discourse_query_engine import prepare_discourse_engine
 
 

diff --git a/tests/unit/test_summary_retriever_base.py b/tests/unit/test_summary_retriever_base.py
@@ -3,8 +3,8 @@
 from unittest.mock import MagicMock
 
 from bot.retrievers.summary_retriever_base import BaseSummarySearch
-from llama_index import Document, MockEmbedding, ServiceContext, VectorStoreIndex
-from llama_index.schema import NodeWithScore
+from llama_index.core import Document, MockEmbedding, Settings, VectorStoreIndex
+from llama_index.core.schema import NodeWithScore
 
 
 class TestSummaryRetrieverBase(TestCase):
@@ -13,11 +13,10 @@ def test_initialize_class(self):
         doc = Document(text="SAMPLESAMPLESAMPLE")
         mock_embedding_model = partial(MockEmbedding, embed_dim=1024)
 
-        service_context = ServiceContext.from_defaults(
-            llm=None, embed_model=mock_embedding_model()
-        )
+        Settings.llm = None
+        Settings.embed_model = mock_embedding_model()
         BaseSummarySearch._setup_index.return_value = VectorStoreIndex.from_documents(
-            documents=[doc], service_context=service_context
+            documents=[doc]
         )
 
         base_summary_search = BaseSummarySearch(

diff --git a/utils/query_engine/level_based_platform_query_engine.py b/utils/query_engine/level_based_platform_query_engine.py
@@ -3,13 +3,16 @@
 from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever
 from bot.retrievers.retrieve_similar_nodes import RetrieveSimilarNodes
 from bot.retrievers.utils.load_hyperparams import load_hyperparams
-from llama_index import VectorStoreIndex
-from llama_index.llms import OpenAI
-from llama_index.prompts import PromptTemplate
-from llama_index.query_engine import CustomQueryEngine
-from llama_index.response_synthesizers import BaseSynthesizer, get_response_synthesizer
-from llama_index.retrievers import BaseRetriever
-from llama_index.schema import NodeWithScore
+from llama_index.core import VectorStoreIndex
+from llama_index.core.prompts import PromptTemplate
+from llama_index.core.query_engine import CustomQueryEngine
+from llama_index.core.response_synthesizers import (
+    BaseSynthesizer,
+    get_response_synthesizer,
+)
+from llama_index.core.retrievers import BaseRetriever
+from llama_index.core.schema import NodeWithScore
+from llama_index.llms.openai import OpenAI
 from tc_hivemind_backend.embeddings.cohere import CohereEmbedding
 from tc_hivemind_backend.pg_vector_access import PGVectorAccess
 from utils.query_engine.level_based_platforms_util import LevelBasedPlatformUtils

diff --git a/utils/query_engine/level_based_platforms_util.py b/utils/query_engine/level_based_platforms_util.py
@@ -1,7 +1,7 @@
 import logging
 
 from dateutil import parser
-from llama_index.schema import NodeWithScore
+from llama_index.core.schema import NodeWithScore
 
 
 class LevelBasedPlatformUtils:

diff --git a/utils/query_engine/prepare_discord_query_engine.py b/utils/query_engine/prepare_discord_query_engine.py
@@ -1,4 +1,4 @@
-from llama_index.query_engine import BaseQueryEngine
+from llama_index.core.query_engine import BaseQueryEngine
 
 from .level_based_platform_query_engine import LevelBasedPlatformQueryEngine
 

diff --git a/utils/query_engine/prepare_discourse_query_engine.py b/utils/query_engine/prepare_discourse_query_engine.py
@@ -1,4 +1,4 @@
-from llama_index.query_engine import BaseQueryEngine
+from llama_index.core.query_engine import BaseQueryEngine
 
 from .level_based_platform_query_engine import LevelBasedPlatformQueryEngine