Aasher commited on
Commit
8fe9ff0
·
1 Parent(s): 03a551b

feat(title-generation): add title generation feature with new model

Browse files

Introduce a new title generation feature using TITLE_GENERATOR_MODEL. Includes:
- Configuration addition for the new model
- Model initialization in llms.py
- New title_generator.py with structured output handling

core/config.py CHANGED
@@ -33,6 +33,7 @@ class Settings(BaseSettings):
33
  FALLBACK_MODEL: str
34
  EMBEDDINGS_MODEL: str
35
  OUTPUT_FORMATTER_MODEL: str
 
36
 
37
  def get_settings():
38
  return Settings()
 
33
  FALLBACK_MODEL: str
34
  EMBEDDINGS_MODEL: str
35
  OUTPUT_FORMATTER_MODEL: str
36
+ TITLE_GENERATOR_MODEL: str
37
 
38
  def get_settings():
39
  return Settings()
workflow/llms.py CHANGED
@@ -15,4 +15,6 @@ main_model = (
15
  output_formatter_model = ChatGroq(
16
  model=settings.OUTPUT_FORMATTER_MODEL,
17
  temperature=0
18
- )
 
 
 
15
  output_formatter_model = ChatGroq(
16
  model=settings.OUTPUT_FORMATTER_MODEL,
17
  temperature=0
18
+ )
19
+
20
+ title_generator_model = init_chat_model(settings.TITLE_GENERATOR_MODEL, temperature=0.2)
workflow/title_generator.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pydantic import BaseModel, Field
2
+ from .llms import title_generator_model
3
+
4
+ from core.config import get_settings
5
+
6
+ settings = get_settings()
7
+
8
+ class Title(BaseModel):
9
+ title: str = Field(..., description="The title based on user message.")
10
+
11
+ async def generate_chat_title(message: str) -> str:
12
+ title = await title_generator_model.with_structured_output(Title).ainvoke(
13
+ "Generate a 4-5 words title based on the following user message. \nUser Message: "
14
+ + message
15
+ )
16
+ return title.title