ai_workflows / app /workflows /til /suggest_headlines_v2.py
theRealNG's picture
restructure the codebase
1f39bb8
from langchain import callbacks, hub
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field, UUID4
from typing import List
import os
import random
class HeadlineInfo(BaseModel):
headline: str
tone: str
reason: str
class Response(BaseModel):
run_id : UUID4
headlines_details: List[HeadlineInfo]
class HeadlineInfoLLM(BaseModel):
headline: str = Field(
description="The suggested headline for the given TIL.",
)
tone: str = Field(
description="The tone of the suggested headline.",
)
clickability_score: int = Field(
description="Evaluate the headline's quality on a scale of 1 to 10 w.r.t the tone.",
)
reason: str = Field(
description="Reason for the clickability_score in one sentence",
)
class HeadlineResults(BaseModel):
headlines_details: List[HeadlineInfoLLM]
class SuggestHeadlinesV2():
def kickoff(self, inputs=[]) -> Response:
self.content = inputs["content"]
return self._get_til_headline()
def _get_til_headline(self) -> Response:
prompt = hub.pull("til_suggest_headline")
llm = ChatOpenAI(model=os.environ['OPENAI_MODEL'], temperature=0.8)
parser = JsonOutputParser(pydantic_object=HeadlineResults)
chain = (prompt | llm | parser).with_config({
"tags": ["til"], "run_name": "Suggest TIL Headlines",
"metadata": {
"version": "v2.0.0",
"growth_activity": "til",
"env": os.environ["ENV"],
"model": os.environ["OPENAI_MODEL"]
}
})
with callbacks.collect_runs() as cb:
self.llm_response = chain.invoke({
"til_content": self.content,
"format_instructions": parser.get_format_instructions(),
})
self.run_id = cb.traced_runs[0].id
return self._handle_response()
def _handle_response(self) -> Response:
response = Response(
run_id=self.run_id,
headlines_details=[]
)
headlines_data = self.llm_response["headlines_details"]
random.shuffle(headlines_data)
for headline_datum in headlines_data[:3]:
response.headlines_details.append(
HeadlineInfo(
headline=headline_datum["headline"],
tone=headline_datum["tone"],
reason=headline_datum["reason"],
)
)
return response