Refactor v1.3 to make it configurable and evaluation pipeline

README.md

@@ -63,6 +63,14 @@ WANDBOT_API_URL="https://localhost:8000"
 WANDB_TRACING_ENABLED="true"
 WANDB_PROJECT="wandbot-dev"
 WANDB_ENTITY="wandbot"
+WANDB_REPORT_API_ENABLE_V2="true"
+WANDB_REPORT_API_DISABLE_MESSAGE="true"
+```
+
+The best practive would be to create a `.env` file in the root of this repo and execute this command in your terminal:
+
+```
+set -o allexport; source .env; set +o allexport
 ```
 
 Once these environment variables are set, you can start the Q&A bot application using the following commands:

pyproject.toml

@@ -12,8 +12,8 @@ include = ["src/**/*", "LICENSE", "README.md"]
 [tool.poetry.dependencies]
 python = ">=3.10.0,<3.12"
 numpy = "^1.26.1"
-wandb = "<=0.16.1"
-tiktoken = "^0.5.1"
+wandb = "<=0.16.3"
+tiktoken = "^0.6.0"
 pandas = "^2.1.2"
 unstructured = "^0.12.3"
 pydantic-settings = "^2.0.3"
@@ -34,19 +34,21 @@ zenpy = "^2.0.46"
 openai = "^1.3.2"
 weave = "^0.31.0"
 colorlog = "^6.8.0"
-litellm = "^1.15.1"
+litellm = "^1.31.6"
 google-cloud-bigquery = "^3.14.1"
 db-dtypes = "^1.2.0"
 python-frontmatter = "^1.1.0"
 pymdown-extensions = "^10.5"
 langchain = "^0.1.5"
-langchain-openai = "^0.0.5"
+langchain-openai = "^0.0.8"
 chromadb = "^0.4.22"
 langchain-experimental = "^0.0.50"
-simsimd = "^3.7.4"
+simsimd = "3.7.7"
 
 [tool.poetry.dev-dependencies]
 #fasttext = {git = "https://github.com/cfculhane/fastText"} # FastText doesn't come with pybind11 and we need to use this workaround.
+llama_index = "^0.10.20"
+ragas = "^0.1.4"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

src/wandbot/api/app.py

@@ -28,6 +28,9 @@
 It uses logger from the utils module for logging purposes.
 """
 
+import os
+os.environ["CONFIG_PATH"] = "config.yaml"
+
 import asyncio
 from contextlib import asynccontextmanager
 from datetime import datetime, timezone
@@ -38,13 +41,21 @@
 from wandbot.api.routers import chat as chat_router
 from wandbot.api.routers import database as database_router
 from wandbot.api.routers import retrieve as retrieve_router
-from wandbot.ingestion.config import VectorStoreConfig
+from wandbot.ingestion.utils import get_embedding_config
 from wandbot.retriever import VectorStore
-from wandbot.utils import get_logger
+from wandbot.utils import get_logger, load_config
 
 logger = get_logger(__name__)
 last_backup = datetime.now().astimezone(timezone.utc)
 
+config = load_config(os.environ["CONFIG_PATH"])
+logger.info(
+ f"Loaded configuration from {os.environ['CONFIG_PATH']}: {config}"
+)
+vectorstore_config = get_embedding_config(
+ config.embeddings.type, config.embeddings.config
+)
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -56,12 +67,18 @@ async def lifespan(app: FastAPI):
  Returns:
  None
  """
- vector_store = VectorStore.from_config(VectorStoreConfig())
- chat_router.chat = chat_router.Chat(vector_store=vector_store)
+ vector_store = VectorStore.from_config(vectorstore_config)
+ chat_router.chat = chat_router.Chat(vector_store=vector_store, config=config)
  database_router.db_client = database_router.DatabaseClient()
- retrieve_router.retriever = retrieve_router.SimpleRetrievalEngine(
- vector_store=vector_store
- )
+ if not config.retrieval_re_ranker.enabled:
+ retrieve_router.retriever = retrieve_router.SimpleRetrievalEngine(
+ vector_store=vector_store
+ )
+ else:
+ retrieve_router.retriever = retrieve_router.SimpleRetrievalEngineWithRerank(
+ vector_store=vector_store
+ )
+
 
  async def backup_db():
  """Periodically backs up the database to a table.

src/wandbot/api/routers/retrieve.py

@@ -4,14 +4,17 @@
 from pydantic import BaseModel
 from starlette import status
 
-from wandbot.retriever.base import SimpleRetrievalEngine
+from wandbot.retriever.base import (
+ SimpleRetrievalEngine,
+ SimpleRetrievalEngineWithRerank,
+)
 
 router = APIRouter(
  prefix="/retrieve",
  tags=["retrievers"],
 )
 
-retriever: SimpleRetrievalEngine | None = None
+retriever: SimpleRetrievalEngineWithRerank | None = None
 
 
 class APIRetrievalResult(BaseModel):
@@ -53,6 +56,7 @@ async def retrieve(request: APIRetrievalRequest) -> APIRetrievalResponse:
  sources=request.sources,
  )
 
+ # TODO: Fix this
  return APIRetrievalResponse(
  query=request.query,
  top_k=[

src/wandbot/chat/chat.py

@@ -32,8 +32,7 @@
 from wandbot.chat.schemas import ChatRequest, ChatResponse
 from wandbot.database.schemas import QuestionAnswer
 from wandbot.retriever import VectorStore
-from wandbot.utils import Timer, get_logger
-from weave.monitoring import StreamTable
+from wandbot.utils import Timer, get_logger, RAGPipelineConfig
 
 logger = get_logger(__name__)
 
@@ -49,29 +48,28 @@ class Chat:
  config: ChatConfig = ChatConfig()
 
  def __init__(
- self, vector_store: VectorStore, config: ChatConfig | None = None
+ self, 
+ vector_store: VectorStore,
+ config: RAGPipelineConfig,
+ chat_config: ChatConfig | None = None
  ):
  """Initializes the Chat instance.
 
  Args:
  config: An instance of ChatConfig containing configuration settings.
  """
  self.vector_store = vector_store
- if config is not None:
- self.config = config
+ self.config = config
+ if chat_config is not None:
+ self.chat_config = chat_config
  self.run = wandb.init(
- project=self.config.wandb_project,
- entity=self.config.wandb_entity,
+ project=self.config.project,
+ entity=self.config.entity,
  job_type="chat",
  )
  self.run._label(repo="wandbot")
- self.chat_table = StreamTable(
- table_name="chat_logs",
- project_name=self.config.wandb_project,
- entity_name=self.config.wandb_entity,
- )
 
- self.rag_pipeline = RAGPipeline(vector_store=vector_store)
+ self.rag_pipeline = RAGPipeline(vector_store=vector_store, config=config)
 
  def _get_answer(
  self, question: str, chat_history: List[QuestionAnswer]
@@ -105,10 +103,10 @@ def __call__(self, chat_request: ChatRequest) -> ChatResponse:
  "total_tokens": result.total_tokens,
  "prompt_tokens": result.prompt_tokens,
  "completion_tokens": result.completion_tokens,
+ "completion_cost": result.completion_cost,
  }
  result_dict.update({"application": chat_request.application})
  self.run.log(usage_stats)
- self.chat_table.log(result_dict)
  return ChatResponse(**result_dict)
  except Exception as e:
  with Timer() as timer:

src/wandbot/chat/config.py

@@ -27,5 +27,3 @@ class ChatConfig(BaseSettings):
  env="WANDB_INDEX_ARTIFACT",
  validation_alias="wandb_index_artifact",
  )
- wandb_project: str | None = Field("wandbot_public", env="WANDB_PROJECT")
- wandb_entity: str | None = Field("wandbot", env="WANDB_ENTITY")

src/wandbot/chat/rag.py

@@ -3,30 +3,22 @@
 
 from langchain_community.callbacks import get_openai_callback
 from pydantic import BaseModel
-from wandbot.rag.query_handler import QueryEnhancer
-from wandbot.rag.response_synthesis import ResponseSynthesizer
-from wandbot.rag.retrieval import FusionRetrieval
-from wandbot.retriever import VectorStore
-from wandbot.utils import Timer, get_logger
 
-logger = get_logger(__name__)
+import litellm
 
+from wandbot.retriever import VectorStore
+from wandbot.utils import Timer, get_logger, RAGPipelineConfig
+from wandbot.rag.query_handler import QueryEnhancer
+from wandbot.rag.retrieval import FusionRetrieval, SimpleRetrievalWithoutFusion
+from wandbot.rag.response_synthesis import ResponseSynthesizer, SimpleResponseSynthesizer
+from wandbot.retriever.base import SimpleRetrievalEngine, SimpleRetrievalEngineWithRerank
+from wandbot.rag.utils import LiteLLMTokenCostLogger
 
-def get_stats_dict_from_token_callback(token_callback):
- return {
- "total_tokens": token_callback.total_tokens,
- "prompt_tokens": token_callback.prompt_tokens,
- "completion_tokens": token_callback.completion_tokens,
- "successful_requests": token_callback.successful_requests,
- }
 
+logger = get_logger(__name__)
 
-def get_stats_dict_from_timer(timer):
- return {
- "start_time": timer.start,
- "end_time": timer.stop,
- "time_taken": timer.elapsed,
- }
+litellm_token_and_cost_counter = LiteLLMTokenCostLogger()
+litellm.callbacks = [litellm_token_and_cost_counter]
 
 
 class RAGPipelineOutput(BaseModel):
@@ -39,6 +31,7 @@ class RAGPipelineOutput(BaseModel):
  total_tokens: int
  prompt_tokens: int
  completion_tokens: int
+ completion_cost: float
  time_taken: float
  start_time: datetime.datetime
  end_time: datetime.datetime
@@ -48,56 +41,98 @@ class RAGPipeline:
  def __init__(
  self,
  vector_store: VectorStore,
+ config: RAGPipelineConfig,
  top_k: int = 5,
  search_type: str = "mmr",
  ):
  self.vector_store = vector_store
- self.query_enhancer = QueryEnhancer()
- self.retrieval = FusionRetrieval(
- vector_store=vector_store, top_k=top_k, search_type=search_type
- )
- self.response_synthesizer = ResponseSynthesizer()
+ self.config = config
+
+ if config.query_enhancer.enabled:
+ self.query_enhancer = QueryEnhancer(config=config)
+ if config.rerank_fusion.enabled:
+ self.retrieval = FusionRetrieval(
+ vector_store=vector_store, top_k=top_k, search_type=search_type
+ )
+ else:
+ self.retrieval = SimpleRetrievalWithoutFusion(
+ vector_store=vector_store, top_k=top_k, search_type=search_type
+ )
+ self.response_synthesizer = ResponseSynthesizer(config=config)
+ else:
+ if config.retrieval_re_ranker.enabled:
+ self.retrieval = SimpleRetrievalEngineWithRerank(
+ vector_store=vector_store, top_k=top_k, search_type=search_type
+ )
+ else:
+ self.retrieval = SimpleRetrievalEngine(
+ vector_store=vector_store, top_k=top_k, search_type=search_type
+ )
+ self.response_synthesizer = SimpleResponseSynthesizer(config=config)
 
  def __call__(
  self, question: str, chat_history: List[Tuple[str, str]] | None = None
- ):
+ ) -> RAGPipelineOutput: 
  if chat_history is None:
  chat_history = []
 
- with get_openai_callback() as query_enhancer_cb, Timer() as query_enhancer_tb:
- enhanced_query = self.query_enhancer.chain.invoke(
- {"query": question, "chat_history": chat_history}
- )
-
- with Timer() as retrieval_tb:
- retrieval_results = self.retrieval.chain.invoke(enhanced_query)
-
- with get_openai_callback() as response_cb, Timer() as response_tb:
+ if self.config.query_enhancer.enabled:
+ with Timer() as query_enhancer_tb:
+ enhanced_query = self.query_enhancer.chain.invoke(
+ {"query": question, "chat_history": chat_history}
+ )
+ with Timer() as retrieval_tb:
+ retrieval_results = self.retrieval.chain.invoke(enhanced_query)
+ else:
+ with Timer() as retrieval_tb:
+ retrieval_results = self.retrieval(
+ question=question, language="en", top_k=self.retrieval.top_k, search_type=self.retrieval.search_type
+ )
+
+ with Timer() as response_tb:
  response = self.response_synthesizer.chain.invoke(retrieval_results)
 
+ question = question if not self.config.query_enhancer.enabled else enhanced_query["standalone_query"]
+ usage_stats = litellm_token_and_cost_counter.get_totals()
+ time_taken = (
+ query_enhancer_tb.elapsed
+ + retrieval_tb.elapsed
+ + response_tb.elapsed
+ if self.config.query_enhancer.enabled
+ else retrieval_tb.elapsed + response_tb.elapsed
+ )
+ start_time = (
+ query_enhancer_tb.start
+ if self.config.query_enhancer.enabled
+ else retrieval_tb.start
+ )
+ end_time = response_tb.stop
+
+ if isinstance(retrieval_results, dict):
+ context = retrieval_results["context"]
+ else:
+ context = retrieval_results.context
+ sources = "\n".join(
+ [
+ item.metadata.get("source")
+ for item in context
+ ]
+ )
+
  output = RAGPipelineOutput(
- question=enhanced_query["standalone_query"],
+ question=question,
  answer=response["response"],
- sources="\n".join(
- [
- item.metadata["source"]
- for item in retrieval_results["context"]
- ]
- ),
+ sources=sources,
  source_documents=response["context_str"],
  system_prompt=response["response_prompt"],
  model=response["response_model"],
- total_tokens=query_enhancer_cb.total_tokens
- + response_cb.total_tokens,
- prompt_tokens=query_enhancer_cb.prompt_tokens
- + response_cb.prompt_tokens,
- completion_tokens=query_enhancer_cb.completion_tokens
- + response_cb.completion_tokens,
- time_taken=query_enhancer_tb.elapsed
- + retrieval_tb.elapsed
- + response_tb.elapsed,
- start_time=query_enhancer_tb.start,
- end_time=response_tb.stop,
+ total_tokens=usage_stats.total_tokens,
+ prompt_tokens=usage_stats.prompt_tokens,
+ completion_tokens=usage_stats.completion_tokens,
+ completion_cost=usage_stats.total_cost,
+ time_taken=time_taken,
+ start_time=start_time,
+ end_time=end_time,
  )
 
  return output