Langgraph Agent - ML4Bio

.gitignore

@@ -1,9 +1,11 @@
 # don't sync coursera docs
 coursera-dl/
 *parsed.json
+*wandb*
+*ipynb*
+*.pem
 
 # don't expose env files
-dummy.ipynb
 .env
 # Created by https://www.toptal.com/developers/gitignore/api/python
 # Edit at https://www.toptal.com/developers/gitignore?templates=python
@@ -37,6 +39,8 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.idea/
+node_modules/
 
 # PyInstaller
 #  Usually these files are written by a python script from a template
@@ -116,7 +120,7 @@ celerybeat.pid
 
 # Environments
 .env
-.venv
+.venv*
 env/
 venv/
 ENV/
@@ -164,3 +168,10 @@ pip-selfcheck.json
 
 # End of https://www.toptal.com/developers/gitignore/api/python
 .aider*
+
+# Input/Output files
+*payload*.json
+*sample*.json
+
+#DS Store
+.DS_Store

.rules/no-bugs.md

@@ -1,8 +1,8 @@
-Make sure our code follows these best practices, UNLESS there's a comment explaining why it's okay to break the rule.
-
-1. Avoid typos.
-2. Don't have like, really obvious bugs.
-3. Don't store secrets in code.
-4. Comment on code smells and bugs, suggest better 'best practices' ways to achieve the same result. No need to be too strict, but please help me learn.
-5. Avoid dangerous stuff, like things that could lead to template injection, SQL injection, broken access control, or really anything that would show up as a CVE somewhere.
-6. If something looks like not best practice, please show how to do it better. We want to learn and improve our code quality and use modern, elegant patterns and practices.
+Make sure our code follows these best practices, UNLESS there's a comment explaining why it's okay to break the rule.
+
+1. Avoid typos.
+2. Don't have like, really obvious bugs.
+3. Don't store secrets in code.
+4. Comment on code smells and bugs, suggest better 'best practices' ways to achieve the same result. No need to be too strict, but please help me learn.
+5. Avoid dangerous stuff, like things that could lead to template injection, SQL injection, broken access control, or really anything that would show up as a CVE somewhere.
+6. If something looks like not best practice, please show how to do it better. We want to learn and improve our code quality and use modern, elegant patterns and practices.

.trunk/.gitignore

@@ -0,0 +1,9 @@
+*out
+*logs
+*actions
+*notifications
+*tools
+plugins
+user_trunk.yaml
+user.yaml
+tmp

.trunk/configs/.isort.cfg

@@ -0,0 +1,2 @@
+[settings]
+profile=google

.trunk/configs/.markdownlint.yaml

@@ -0,0 +1,10 @@
+# Autoformatter friendly markdownlint config (all formatting rules disabled)
+default: true
+blank_lines: false
+bullet: false
+html: false
+indentation: false
+line_length: false
+spaces: false
+url: false
+whitespace: false

.trunk/configs/.shellcheckrc

@@ -0,0 +1,7 @@
+enable=all
+source-path=SCRIPTDIR
+disable=SC2154
+
+# If you're having issues with shellcheck following source, disable the errors via:
+# disable=SC1090
+# disable=SC1091

.trunk/configs/.style.yapf

@@ -0,0 +1,4 @@
+[style]
+based_on_style = google
+column_limit = 140
+indent_width = 2

.trunk/configs/.yamllint.yaml

@@ -0,0 +1,10 @@
+rules:
+  quoted-strings:
+    required: only-when-needed
+    extra-allowed: ["{|}"]
+  empty-values:
+    forbid-in-block-mappings: true
+    forbid-in-flow-mappings: true
+  key-duplicates: {}
+  octal-values:
+    forbid-implicit-octal: true

.trunk/configs/ruff.toml

@@ -0,0 +1,5 @@
+# Generic, formatter-friendly config.
+select = ["B", "D3", "E", "F"]
+
+# Never enforce `E501` (line length violations). This should be handled by formatters.
+ignore = ["E501"]

.trunk/trunk.yaml

@@ -0,0 +1,55 @@
+# This file controls the behavior of Trunk: https://docs.trunk.io/cli
+# To learn more about the format of this file, see https://docs.trunk.io/reference/trunk-yaml
+version: 0.1
+cli:
+  version: 1.22.2
+# Trunk provides extensibility via plugins. (https://docs.trunk.io/plugins)
+plugins:
+  sources:
+    - id: trunk
+      ref: v1.6.0
+      uri: https://github.com/trunk-io/plugins
+# Many linters and tools depend on runtimes - configure them here. (https://docs.trunk.io/runtimes)
+runtimes:
+  enabled:
+    - [email protected]
+    - [email protected]
+    - [email protected]
+# This is the section where you manage your linters. (https://docs.trunk.io/check/configuration)
+# - [email protected]  # too sensitive, causing failures that make devs skip checks.
+lint:
+  enabled:
+    # - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    # - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - git-diff-check
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+    - [email protected]
+  ignore:
+    - linters: [ALL]
+      paths:
+        - .github/**/*
+        - .trunk/**/*
+        - .rules/**/*
+        - mkdocs.yml
+        - .DS_Store
+        - .vscode/**/*
+        - README.md
+actions:
+  enabled:
+    - trunk-announce
+    - trunk-check-pre-push
+    - trunk-fmt-pre-commit
+    - trunk-upgrade-available

ai_ta_backend/agents/MemoryCallbacks.py

@@ -0,0 +1,72 @@
+import os
+from typing import Any, Dict, Union
+
+from langchain.agents import AgentType
+from langchain.agents import initialize_agent
+from langchain.agents import load_tools
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.llms import OpenAI
+from langchain.schema import AgentAction
+from langchain.schema import AgentFinish
+import supabase
+
+
+class MemoryCallbackHandler(BaseCallbackHandler):
+
+  def __init__(self):
+    self.tool_in_progress = False  # usage TBD
+
+    self.supabase_client = supabase.create_client(  # type: ignore
+        supabase_url=os.getenv('SUPABASE_URL'),  # type: ignore
+        supabase_key=os.getenv('SUPABASE_API_KEY'))  # type: ignore
+
+  def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
+    print(f"on_tool_start {serialized}")
+    self.tool_in_progress = True
+
+  def on_tool_end(self, output: str, **kwargs: Any) -> Any:
+    """Run when LLM errors."""
+    print(f"On tool end: {output}")
+    if self.tool_in_progress:
+      self.tool_in_progress = False
+      print(f"Tool output: {output}")
+
+  def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    """Run when LLM errors."""
+    pass
+
+  def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    """Run when LLM errors."""
+    pass
+
+  def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+    print(f"on_agent_action {action}")
+
+  def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> Any:
+    print(f"on_agent_finish {finish}")
+
+  # def on_llm_start(
+  #     self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+  # ) -> Any:
+  #     print(f"on_llm_start {serialized}")
+
+  # def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
+  #     print(f"on_new_token {token}")
+
+  # def on_chain_start(
+  #     self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+  # ) -> Any:
+  #     print(f"on_chain_start {serialized['name']}")
+
+
+if __name__ == "__main__":
+  # just for testing
+  handler1 = MemoryCallbackHandler()
+
+  llm = OpenAI(temperature=0, streaming=True)
+  tools = load_tools(["llm-math"], llm=llm)
+  agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)
+
+  # Callbacks for handler1 will be issued by every object involved in the
+  # Agent execution (llm, llmchain, tool, agent executor)
+  agent.run("What is 2 raised to the 0.235 power?", callbacks=[handler1])