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()