Call get_chat_history_manager() at runtime

Author: mliu-cloudera

llm-service/app/routers/index/sessions/__init__.py

@@ -59,7 +59,7 @@
 from ....services.chat.suggested_questions import generate_suggested_questions
 from ....services.chat_history.chat_history_manager import (
     RagStudioChatMessage,
-    chat_history_manager,
+    get_chat_history_manager,
 )
 from ....services.chat_history.paginator import paginate
 from ....services.metadata_apis import session_metadata_api
@@ -142,7 +142,7 @@ class RagStudioChatHistoryResponse(BaseModel):
 def chat_history(
     session_id: int, limit: Optional[int] = None, offset: Optional[int] = None
 ) -> RagStudioChatHistoryResponse:
-    results = chat_history_manager.retrieve_chat_history(session_id=session_id)
+    results = get_chat_history_manager().retrieve_chat_history(session_id=session_id)
 
     paginated_results, previous_id, next_id = paginate(results, limit, offset)
     return RagStudioChatHistoryResponse(
@@ -158,8 +158,8 @@ def chat_history(
 )
 @exceptions.propagates
 def get_message_by_id(session_id: int, message_id: str) -> RagStudioChatMessage:
-    results: list[RagStudioChatMessage] = chat_history_manager.retrieve_chat_history(
-        session_id=session_id
+    results: list[RagStudioChatMessage] = (
+        get_chat_history_manager().retrieve_chat_history(session_id=session_id)
     )
     for message in results:
         if message.id == message_id:
@@ -175,14 +175,14 @@ def get_message_by_id(session_id: int, message_id: str) -> RagStudioChatMessage:
 )
 @exceptions.propagates
 def clear_chat_history(session_id: int) -> str:
-    chat_history_manager.clear_chat_history(session_id=session_id)
+    get_chat_history_manager().clear_chat_history(session_id=session_id)
     return "Chat history cleared."
 
 
 @router.delete("", summary="Deletes the requested session.")
 @exceptions.propagates
 def delete_session(session_id: int) -> str:
-    chat_history_manager.delete_chat_history(session_id=session_id)
+    get_chat_history_manager().delete_chat_history(session_id=session_id)
     return "Chat history deleted."
 
 

llm-service/app/services/chat/chat.py

@@ -50,7 +50,7 @@
     Evaluation,
     RagMessage,
     RagStudioChatMessage,
-    chat_history_manager,
+    get_chat_history_manager,
 )
 from app.services.metadata_apis.session_metadata_api import Session
 from app.services.mlflow import record_rag_mlflow_run, record_direct_llm_mlflow_run
@@ -172,7 +172,7 @@ def finalize_response(
     record_rag_mlflow_run(
         new_chat_message, query_configuration, response_id, session, user_name
     )
-    chat_history_manager.append_to_history(session.id, [new_chat_message])
+    get_chat_history_manager().append_to_history(session.id, [new_chat_message])
 
     return new_chat_message
 
@@ -198,5 +198,5 @@ def direct_llm_chat(
         timestamp=time.time(),
         condensed_question=None,
     )
-    chat_history_manager.append_to_history(session.id, [new_chat_message])
+    get_chat_history_manager().append_to_history(session.id, [new_chat_message])
     return new_chat_message

llm-service/app/services/chat/streaming_chat.py

@@ -53,7 +53,7 @@
 from app.services.chat_history.chat_history_manager import (
     RagStudioChatMessage,
     RagMessage,
-    chat_history_manager,
+    get_chat_history_manager,
 )
 from app.services.metadata_apis.session_metadata_api import Session
 from app.services.mlflow import record_direct_llm_mlflow_run
@@ -217,4 +217,4 @@ def _stream_direct_llm_chat(
         timestamp=time.time(),
         condensed_question=None,
     )
-    chat_history_manager.append_to_history(session.id, [new_chat_message])
+    get_chat_history_manager().append_to_history(session.id, [new_chat_message])

llm-service/app/services/chat/utils.py

@@ -43,7 +43,7 @@
 from pydantic import BaseModel
 
 from app.services.chat_history.chat_history_manager import (
-    chat_history_manager,
+    get_chat_history_manager,
     RagPredictSourceNode,
 )
 
@@ -54,7 +54,7 @@ class RagContext(BaseModel):
 
 
 def retrieve_chat_history(session_id: int) -> List[RagContext]:
-    chat_history = chat_history_manager.retrieve_chat_history(session_id)[-10:]
+    chat_history = get_chat_history_manager().retrieve_chat_history(session_id)[-10:]
     history: List[RagContext] = []
     for message in chat_history:
         history.append(

llm-service/app/services/chat_history/chat_history_manager.py

@@ -27,7 +27,7 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
-
+import functools
 from abc import ABCMeta, abstractmethod
 from typing import Optional, Literal
 
@@ -85,7 +85,8 @@ def append_to_history(
         pass
 
 
-def _create_chat_history_manager() -> ChatHistoryManager:
+@functools.cache
+def get_chat_history_manager() -> ChatHistoryManager:
     from app.services.chat_history.simple_chat_history_manager import (
         SimpleChatHistoryManager,
     )
@@ -99,6 +100,3 @@ def _create_chat_history_manager() -> ChatHistoryManager:
         return S3ChatHistoryManager()
     else:
         return SimpleChatHistoryManager()
-
-
-chat_history_manager = _create_chat_history_manager()

llm-service/app/services/llm_completion.py

@@ -47,7 +47,7 @@
 from . import models
 from .chat_history.chat_history_manager import (
     RagStudioChatMessage,
-    chat_history_manager,
+    get_chat_history_manager,
 )
 from .query.query_configuration import QueryConfiguration
 
@@ -60,7 +60,7 @@ def make_chat_messages(x: RagStudioChatMessage) -> list[ChatMessage]:
 
 def completion(session_id: int, question: str, model_name: str) -> ChatResponse:
     model = models.LLM.get(model_name)
-    chat_history = chat_history_manager.retrieve_chat_history(session_id)[:10]
+    chat_history = get_chat_history_manager().retrieve_chat_history(session_id)[:10]
     messages = list(
         itertools.chain.from_iterable(
             map(lambda x: make_chat_messages(x), chat_history)
@@ -78,7 +78,7 @@ def stream_completion(
     Returns a generator that yields ChatResponse objects as they become available.
     """
     model = models.LLM.get(model_name)
-    chat_history = chat_history_manager.retrieve_chat_history(session_id)[:10]
+    chat_history = get_chat_history_manager().retrieve_chat_history(session_id)[:10]
     messages = list(
         itertools.chain.from_iterable(
             map(lambda x: make_chat_messages(x), chat_history)

llm-service/app/services/session.py

@@ -40,7 +40,7 @@
 
 from . import models
 from .chat_history.chat_history_manager import (
-    chat_history_manager,
+    get_chat_history_manager,
     RagStudioChatMessage,
 )
 from .metadata_apis import session_metadata_api
@@ -87,7 +87,7 @@
 
 def rename_session(session_id: int, user_name: Optional[str]) -> str:
     chat_history: list[RagStudioChatMessage] = (
-        chat_history_manager.retrieve_chat_history(session_id=session_id)
+        get_chat_history_manager().retrieve_chat_history(session_id=session_id)
     )
     if not chat_history:
         logger.info("No chat history found for session ID %s", session_id)