parameterize research assistant llms

langchain-streamlit-demo/app.py

@@ -26,7 +26,7 @@ from llm_resources import (
  get_runnable,
  get_texts_and_multiretriever,
 )
-from research_assistant.chain import chain as research_assistant_chain
+from research_assistant.chain import get_chain as get_research_assistant_chain
 
 __version__ = "2.0.1"
 
@@ -367,7 +367,7 @@ with sidebar:
 
 
 # --- LLM Instantiation ---
-st.session_state.llm = get_llm(
+get_llm_args = dict(
  provider=st.session_state.provider,
  model=model,
  provider_api_key=provider_api_key,
@@ -382,6 +382,8 @@ st.session_state.llm = get_llm(
  "AZURE_OPENAI_MODEL_VERSION": st.session_state.AZURE_OPENAI_MODEL_VERSION,
  },
 )
+get_llm_args_temp_zero = get_llm_args | {"temperature": 0.0}
+st.session_state.llm = get_llm(**get_llm_args)
 
 # --- Chat History ---
 for msg in STMEMORY.messages:
@@ -448,12 +450,16 @@ if st.session_state.llm:
  WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
  ]
  if st.session_state.provider in ("Azure OpenAI", "OpenAI"):
+ research_assistant_chain = get_research_assistant_chain(
+ search_llm=get_llm(**get_llm_args_temp_zero), # type: ignore
+ writer_llm=get_llm(**get_llm_args_temp_zero), # type: ignore
+ )
  st_callback = StreamlitCallbackHandler(st.container())
  callbacks.append(st_callback)
  research_assistant_tool = Tool.from_function(
  func=lambda s: research_assistant_chain.invoke(
  {"question": s},
- config=get_config(callbacks),
+ # config=get_config(callbacks),
  ),
  name="web-research-assistant",
  description="this assistant returns a comprehensive report based on web research. for quick facts, use duckduckgo instead.",
@@ -473,7 +479,7 @@ if st.session_state.llm:
  doc_chain_tool = Tool.from_function(
  func=lambda s: st.session_state.doc_chain.invoke(
  s,
- config=get_config(callbacks),
+ # config=get_config(callbacks),
  ),
  name="user-document-chat",
  description="this assistant returns a response based on the user's custom context. if the user's meaning is unclear, perhaps the answer is here. generally speaking, try this tool before conducting web research.",

langchain-streamlit-demo/research_assistant/__init__.py

@@ -1,3 +1,3 @@
-from research_assistant.chain import chain
+from research_assistant.chain import get_chain
 
-__all__ = ["chain"]
+__all__ = ["get_chain"]

langchain-streamlit-demo/research_assistant/chain.py

@@ -1,16 +1,18 @@
 from langchain_core.pydantic_v1 import BaseModel
 from langchain_core.runnables import RunnablePassthrough
 
-from research_assistant.search.web import chain as search_chain
-from research_assistant.writer import chain as writer_chain
+from research_assistant.search.web import get_search_chain
+from research_assistant.writer import get_writer_chain
+from langchain.llms.base import BaseLLM
+from langchain.schema.runnable import Runnable
 
-chain_notypes = (
- RunnablePassthrough().assign(research_summary=search_chain) | writer_chain
-)
 
+def get_chain(search_llm: BaseLLM, writer_llm: BaseLLM) -> Runnable:
+ chain_notypes = RunnablePassthrough().assign(
+ research_summary=get_search_chain(search_llm),
+ ) | get_writer_chain(writer_llm)
 
-class InputType(BaseModel):
- question: str
+ class InputType(BaseModel):
+ question: str
 
-
-chain = chain_notypes.with_types(input_type=InputType)
+ return chain_notypes.with_types(input_type=InputType)

langchain-streamlit-demo/research_assistant/search/web.py

@@ -3,7 +3,7 @@ from typing import Any
 
 import requests
 from bs4 import BeautifulSoup
-from langchain.chat_models import ChatOpenAI
+from langchain.llms.base import BaseLLM
 from langchain.prompts import ChatPromptTemplate
 from langchain.retrievers.tavily_search_api import TavilySearchAPIRetriever
 from langchain.utilities import DuckDuckGoSearchAPIWrapper
@@ -130,25 +130,6 @@ Using the above text, answer in short the following question:
 if the question cannot be answered using the text, imply summarize the text. Include all factual information, numbers, stats etc if available.""" # noqa: E501
 SUMMARY_PROMPT = ChatPromptTemplate.from_template(SUMMARY_TEMPLATE)
 
-scrape_and_summarize: Runnable[Any, Any] = (
- RunnableParallel(
- {
- "question": lambda x: x["question"],
- "text": lambda x: scrape_text(x["url"])[:10000],
- "url": lambda x: x["url"],
- },
- )
- | RunnableParallel(
- {
- "summary": SUMMARY_PROMPT | ChatOpenAI(temperature=0) | StrOutputParser(),
- "url": lambda x: x["url"],
- },
- )
- | RunnableLambda(lambda x: f"Source Url: {x['url']}\nSummary: {x['summary']}")
-)
-
-multi_search = get_links | scrape_and_summarize.map() | (lambda x: "\n".join(x))
-
 
 def load_json(s):
  try:
@@ -157,24 +138,41 @@ def load_json(s):
  return {}
 
 
-search_query = SEARCH_PROMPT | ChatOpenAI(temperature=0) | StrOutputParser() | load_json
-choose_agent = (
- CHOOSE_AGENT_PROMPT | ChatOpenAI(temperature=0) | StrOutputParser() | load_json
-)
-
-get_search_queries = (
- RunnablePassthrough().assign(
- agent_prompt=RunnableParallel({"task": lambda x: x})
- | choose_agent
- | (lambda x: x.get("agent_role_prompt")),
+def get_search_chain(model: BaseLLM) -> Runnable:
+ scrape_and_summarize: Runnable[Any, Any] = (
+ RunnableParallel(
+ {
+ "question": lambda x: x["question"],
+ "text": lambda x: scrape_text(x["url"])[:10000],
+ "url": lambda x: x["url"],
+ },
+ )
+ | RunnableParallel(
+ {
+ "summary": SUMMARY_PROMPT | model | StrOutputParser(),
+ "url": lambda x: x["url"],
+ },
+ )
+ | RunnableLambda(lambda x: f"Source Url: {x['url']}\nSummary: {x['summary']}")
  )
- | search_query
-)
 
+ multi_search = get_links | scrape_and_summarize.map() | (lambda x: "\n".join(x))
 
-chain = (
- get_search_queries
- | (lambda x: [{"question": q} for q in x])
- | multi_search.map()
- | (lambda x: "\n\n".join(x))
-)
+ search_query = SEARCH_PROMPT | model | StrOutputParser() | load_json
+ choose_agent = CHOOSE_AGENT_PROMPT | model | StrOutputParser() | load_json
+
+ get_search_queries = (
+ RunnablePassthrough().assign(
+ agent_prompt=RunnableParallel({"task": lambda x: x})
+ | choose_agent
+ | (lambda x: x.get("agent_role_prompt")),
+ )
+ | search_query
+ )
+
+ return (
+ get_search_queries
+ | (lambda x: [{"question": q} for q in x])
+ | multi_search.map()
+ | (lambda x: "\n\n".join(x))
+ )

langchain-streamlit-demo/research_assistant/writer.py

@@ -1,7 +1,8 @@
-from langchain.chat_models import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import ConfigurableField
+from langchain.llms.base import BaseLLM
+from langchain.schema.runnable import Runnable
 
 WRITER_SYSTEM_PROMPT = "You are an AI critical thinker research assistant. Your sole purpose is to write well written, critically acclaimed, objective and structured reports on given text." # noqa: E501
 
@@ -50,7 +51,6 @@ Use appropriate Markdown syntax to format the outline and ensure readability.
 
 Please do your best, this is very important to my career.""" # noqa: E501
 
-model = ChatOpenAI(temperature=0)
 prompt = ChatPromptTemplate.from_messages(
  [
  ("system", WRITER_SYSTEM_PROMPT),
@@ -72,4 +72,7 @@ prompt = ChatPromptTemplate.from_messages(
  ],
  ),
 )
-chain = prompt | model | StrOutputParser()
+
+
+def get_writer_chain(model: BaseLLM) -> Runnable:
+ return prompt | model | StrOutputParser()