Coverage for vorlagellm/chains.py: 82.14%

1from langchain.schema.output_parser import StrOutputParser 

2import re 

3 

4 

5from .prompts import build_prompt, build_source_prompt, build_corresponding_text_prompt 

6 

7 

8def parse_result(output:str) -> tuple[str,str]: 

9 components = re.split(r'---+', output) 

10 if len(components) < 2: 

11 readings = output 

12 justification = output 

13 else: 

14 readings = components[0] 

15 justification = "".join(components[1:]).strip() 

16 

17 readings = [int(i)-1 for i in re.findall(r'\d+', readings)] # readings listed from 1 onwards 

18 return readings, justification 

19 

20 

21def strip_hyphens(text:str) -> str: 

22 return re.sub(r'---+.*', '', text).strip() 

23 

24 

25 

26def build_chain(llm, doc_language: str, apparatus_language: str, initiate_response:bool=False): 

27 prompt = build_prompt(doc_language=doc_language, apparatus_language=apparatus_language, initiate_response=initiate_response) 

28 

29 return prompt | llm | StrOutputParser() | parse_result 

30 

31 

32def build_source_chain(llm, doc_language: str, apparatus_language: str, notes:str, initiate_response:bool=False): 

33 prompt = build_source_prompt(doc_language=doc_language, apparatus_language=apparatus_language, notes=notes, initiate_response=initiate_response) 

34 

35 return prompt | llm | StrOutputParser() | parse_result 

36 

37 

38def print_prompt(prompt): 

39 print(prompt.to_string()) 

40 return prompt 

41 

42 

43def build_corresponding_text_chain(llm, doc_language: str, apparatus_language: str, verbose:bool=False, initiate_response:bool=False): 

44 prompt = build_corresponding_text_prompt(doc_language=doc_language, apparatus_language=apparatus_language, initiate_response=initiate_response) 

45 if verbose: 

46 prompt = prompt | print_prompt 

47 

48 # llm_with_fallback = llm.bind(stop=["----"]).with_fallbacks([llm]) 

49 

50 return prompt | llm | StrOutputParser() | strip_hyphens