Some more refactoring

langsim/__init__.py

@@ -1,11 +1,8 @@
-from langsim.shell import dialog as dialog_python
-
-
 def load_ipython_extension(ipython):
     """
     Any module file that define a function named `load_ipython_extension`
     can be loaded via `%load_ext module.path` or be configured to be
     autoloaded by IPython at startup time.
     """
-    from .magics import CompMatMagics
+    from langsim.magics import CompMatMagics
     ipython.register_magics(CompMatMagics)

langsim/llm.py

@@ -1,49 +1,45 @@
 from langchain.agents import AgentExecutor
-from langchain_openai import ChatOpenAI
+from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.agents.format_scratchpad.openai_tools import (
     format_to_openai_tool_messages,
 )
-from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser
-
-
-SYSTEM_PROMPT = """
-You are very powerful assistant, but don't know current events.
-You can calculate materials properties with EMT and MACE.
-- Effective medium theory (EMT) is a cheap, analytical model that describes the macroscopic properties of composite materials, but is not available for all elements. 
-  To calculate with EMT use the calculator string "emt".
-- MACE is a powerful machine learning force field for predicting many-body atomic interactions that covers the periodic table.
-  To calculate with MACE use the calculator string "mace".
-If the user does not specify a calculator string, ask the user to provide one.
 
-If no chemical element is provided, use Aluminum as the default chemical element.
-"""
-# and emt as the default calculator string
+from langsim.tools.interface import (
+    get_bulk_modulus,
+    get_equilibrium_volume,
+    get_experimental_elastic_property_wikipedia,
+    plot_equation_of_state,
+    get_equilibrium_lattice,
+)
+from langsim.prompt import SYSTEM_PROMPT
 
-def get_executor(OPENAI_API_KEY):
-    from langsim.tools import interface
 
-    llm = ChatOpenAI(
-        model="gpt-4", temperature=0, openai_api_key=OPENAI_API_KEY
-        # model="gpt-3.5-turbo", temperature=0, openai_api_key=OPENAI_API_KEY
-    )
+def get_executor(api_provider, api_key, api_url=None, api_model=None, api_temperature=0):
+    if api_provider.lower() == "openai":
+        print("loading openai")
+        from langchain_openai import ChatOpenAI
+        if api_model is None:
+            api_model = "gpt-4"
+        llm = ChatOpenAI(
+            model=api_model, temperature=api_temperature, openai_api_key=api_key, base_url=api_url,
+        )
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", SYSTEM_PROMPT),
+                ("placeholder", "{conversation}"),
+                MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ]
+        )
+    else:
+        raise ValueError()
     tools = [
-        interface.get_equilibrium_volume,
-        interface.get_equilibrium_lattice,
-        interface.plot_equation_of_state,
-        interface.get_bulk_modulus,
-        interface.get_experimental_elastic_property_wikipedia,
+        get_equilibrium_volume,
+        get_equilibrium_lattice,
+        plot_equation_of_state,
+        get_bulk_modulus,
+        get_experimental_elastic_property_wikipedia,
     ]
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            (
-                "system",
-                SYSTEM_PROMPT,
-            ),
-            ("placeholder", "{conversation}"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
-        ]
-    )
     llm_with_tools = llm.bind_tools(tools)
     agent = (
         {

langsim/magics.py

@@ -7,27 +7,40 @@
 # Stdlib imports
 import os
 
-# from IPython/Jupyter APIs
-from IPython.core.magic import (Magics, magics_class, line_cell_magic)
-
-from IPython.core.magic_arguments import (magic_arguments, argument,
-parse_argstring)
-
+from IPython import get_ipython
+from IPython.core.magic import (
+    Magics,
+    magics_class,
+    line_cell_magic,
+)
+from IPython.core.magic_arguments import (
+    magic_arguments,
+    argument,
+    parse_argstring,
+)
 from IPython.display import Markdown
-from .llm import get_executor
+from langsim.llm import get_executor
+
 
 def get_output(messages):
-    agent_executor = get_executor(OPENAI_API_KEY=os.getenv("OPENAI_API_KEY"))
+    env = os.environ
+    agent_executor = get_executor(
+        api_provider=env.get("LANGSIM_PROVIDER", "OPENAI"),
+        api_key=env.get("LANGSIM_API_KEY"),
+        api_url=env.get("LANGSIM_API_URL", None),
+        api_model=env.get("LANGSIM_MODEL", None),
+        api_temperature=env.get("LANGSIM_TEMP", 0),
+    )
     return list(agent_executor.stream({"conversation": messages}))[-1]
 
+
 # Class to manage state and expose the main magics
 # The class MUST call this class decorator at creation time
 @magics_class
 class CompMatMagics(Magics):
     def __init__(self, shell):
         # You must call the parent constructor
         super(CompMatMagics, self).__init__(shell)
-        self.api_key = os.getenv("OPENAI_API_KEY")
         self.last_code = ""
         self.messages = []
 
@@ -47,7 +60,9 @@ def __init__(self, shell):
         """)
     @line_cell_magic
     def chat(self, line, cell=None):
-        "Chat with GPTChat."
+        """
+        Chat with GPTChat.
+        """
         args = parse_argstring(self.chat, line)
 
         if cell is None:
@@ -58,14 +73,12 @@ def chat(self, line, cell=None):
         response = get_output(self.messages)
         output = response['output']
         self.messages.append(("ai", output))
-        #output = response.choices[0].text.strip()
         if args.raw:
             return Markdown(f"```\n{output}\n```\n")
         else:
             return Markdown(output)
 
 
-
 # If testing interactively, it's convenient to %run as a script in Jupyter
 if __name__ == "__main__":
     get_ipython().register_magics(CompMatMagics)

langsim/prompt.py

@@ -0,0 +1,11 @@
+SYSTEM_PROMPT = """
+You are very powerful assistant, but don't know current events.
+You can calculate materials properties with EMT and MACE.
+- Effective medium theory (EMT) is a cheap, analytical model that describes the macroscopic properties of composite materials, but is not available for all elements. 
+  To calculate with EMT use the calculator string "emt".
+- MACE is a powerful machine learning force field for predicting many-body atomic interactions that covers the periodic table.
+  To calculate with MACE use the calculator string "mace".
+If the user does not specify a calculator string, ask the user to provide one.
+
+If no chemical element is provided, use Aluminum as the default chemical element.
+"""