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: Added update based on new lib changes! #39

Merged
merged 9 commits into from
Mar 13, 2024
4 changes: 2 additions & 2 deletions bot/retrievers/forum_summary_retriever.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
9 changes: 5 additions & 4 deletions bot/retrievers/retrieve_similar_nodes.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
7 changes: 3 additions & 4 deletions bot/retrievers/summary_retriever_base.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
4 changes: 2 additions & 2 deletions discord_query.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
2 changes: 1 addition & 1 deletion docker-entrypoint.sh
Original file line number Diff line number Diff line change
@@ -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
5 changes: 3 additions & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
numpy
llama-index==0.9.48
pymongo
python-dotenv
pgvector
Expand All @@ -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
Expand Down
21 changes: 12 additions & 9 deletions subquery.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -97,17 +98,19 @@ def query_multiple_source(
if github:
raise NotImplementedError

embed_model = CohereEmbedding()
llm = OpenAI("gpt-4")
cyri113 marked this conversation as resolved.
Show resolved Hide resolved
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))
Expand Down
2 changes: 1 addition & 1 deletion tests/integration/test_retrieve_similar_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
11 changes: 4 additions & 7 deletions tests/unit/test_discord_summary_retriever.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_level_based_platform_query_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_level_based_platform_util.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_prepare_discord_query_engine.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
2 changes: 1 addition & 1 deletion tests/unit/test_prepare_discourse_query_engine.py
Original file line number Diff line number Diff line change
@@ -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


Expand Down
11 changes: 5 additions & 6 deletions tests/unit/test_summary_retriever_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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(
Expand Down
17 changes: 10 additions & 7 deletions utils/query_engine/level_based_platform_query_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion utils/query_engine/level_based_platforms_util.py
Original file line number Diff line number Diff line change
@@ -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:
Expand Down
2 changes: 1 addition & 1 deletion utils/query_engine/prepare_discord_query_engine.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion utils/query_engine/prepare_discourse_query_engine.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
Loading