From 4614dbb237e9984cd5b45e24aba495c11e65978d Mon Sep 17 00:00:00 2001
From: James Ketrenos <james_git@ketrenos.com>
Date: Tue, 29 Apr 2025 17:46:10 -0700
Subject: [PATCH] Almost working?

---
 Dockerfile               |  2 +-
 src/server.py            |  3 +++
 src/utils/__init__.py    | 48 +++++++++++++++++++++++++++++++++++++++-
 src/utils/agents/base.py |  9 ++++++--
 src/utils/context.py     |  4 +++-
 5 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index edc1aa3..8c04ffc 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -293,7 +293,7 @@ RUN { \
     echo '    openssl req -x509 -nodes -days 365 -newkey rsa:2048 -keyout src/key.pem -out src/cert.pem -subj "/C=US/ST=OR/L=Portland/O=Development/CN=localhost"'; \
     echo '  fi' ; \
     echo '  while true; do'; \
-    echo '    if [[ ! -e /opt/backstory/block-server ]]; then' \
+    echo '    if [[ ! -e /opt/backstory/block-server ]]; then'; \
     echo '      echo "Launching Backstory server..."'; \
     echo '      python src/server.py "${@}" || echo "Backstory server died."'; \
     echo '    else'; \
diff --git a/src/server.py b/src/server.py
index 61ab2f0..2a91789 100644
--- a/src/server.py
+++ b/src/server.py
@@ -1,3 +1,6 @@
+import warnings
+warnings.filterwarnings("ignore", message="Overriding a previously registered kernel")
+
 # %%
 # Imports [standard]
 # Standard library modules (no try-except needed)
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index e00511e..eea4e13 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,9 +1,12 @@
+from typing import Optional, Type
+
 from . import defines  
 from .rag import ChromaDBFileWatcher, start_file_watcher
 from .message import Message
 from .conversation import Conversation
 from .context import Context
 from . import agents
+import logging
 
 from .agents import Agent, __all__ as agents_all
 
@@ -14,4 +17,47 @@ __all__ = [
   'Message', 
   'ChromaDBFileWatcher', 
   'start_file_watcher'
-] + agents_all
\ No newline at end of file
+] + agents_all
+
+# Resolve circular dependencies by rebuilding models
+# Call model_rebuild() on Agent and Context
+Agent.model_rebuild()
+Context.model_rebuild()
+import logging
+import importlib
+from pydantic import BaseModel
+from typing import Type
+
+# Assuming class_registry is available from agents/__init__.py
+from .agents import class_registry, AnyAgent
+
+def rebuild_models():
+    Context.model_rebuild()
+    for class_name, (module_name, _) in class_registry.items():
+        try:
+            module = importlib.import_module(module_name)
+            cls = getattr(module, class_name, None)
+            
+            logging.info(f"Checking: {class_name} in module {module_name}")
+            logging.info(f"  cls: {True if cls else False}")
+            logging.info(f"  isinstance(cls, type): {isinstance(cls, type)}")
+            logging.info(f"  issubclass(cls, BaseModel): {issubclass(cls, BaseModel) if cls else False}")
+            logging.info(f"  issubclass(cls, AnyAgent): {issubclass(cls, AnyAgent) if cls else False}")
+            logging.info(f"  cls is not AnyAgent: {cls is not AnyAgent if cls else True}")
+
+            if (
+                cls
+                and isinstance(cls, type)
+                and issubclass(cls, BaseModel)
+                and issubclass(cls, AnyAgent)
+                and cls is not AnyAgent
+            ):
+                logging.info(f"Rebuilding {class_name} from {module_name}")
+                cls.model_rebuild()
+        except ImportError as e:
+            logging.error(f"Failed to import module {module_name}: {e}")
+        except Exception as e:
+            logging.error(f"Error processing {class_name} in {module_name}: {e}")
+
+# Call this after all modules are imported
+rebuild_models()
\ No newline at end of file
diff --git a/src/utils/agents/base.py b/src/utils/agents/base.py
index f77deb3..7e06312 100644
--- a/src/utils/agents/base.py
+++ b/src/utils/agents/base.py
@@ -1,5 +1,5 @@
 from pydantic import BaseModel, Field, model_validator, PrivateAttr
-from typing import Literal, TypeAlias, get_args, List, Generator, Iterator, AsyncGenerator, TYPE_CHECKING, Optional, ClassVar
+from typing import Literal, TypeAlias, get_args, List, Generator, Iterator, AsyncGenerator, TYPE_CHECKING, Optional, ClassVar, ForwardRef
 from abc import ABC, abstractmethod
 from typing_extensions import Annotated
 import logging
@@ -8,7 +8,9 @@ import logging
 if TYPE_CHECKING:
   from .. context import Context
 
-from .types import AgentBase, ContextRef, registry
+ContextRef = ForwardRef('Context')
+
+from .types import AgentBase, registry
 
 from .. conversation import Conversation
 from .. message import Message
@@ -41,6 +43,9 @@ class Agent(AgentBase):
         # Set agent_type from class if not provided
         if 'agent_type' not in data:
             data['agent_type'] = self.__class__.agent_type
+        from .. context import Context
+        Context.model_rebuild()
+        self.__class__.model_rebuild()
         super().__init__(**data)
 
     def get_agent_type(self):
diff --git a/src/utils/context.py b/src/utils/context.py
index a9582ab..da5248d 100644
--- a/src/utils/context.py
+++ b/src/utils/context.py
@@ -176,4 +176,6 @@ class Context(BaseModel):
                 summary += f"\nFacts: {agent.facts}\n"
             elif agent.agent_type == "chat":
                 summary += f"\nChat Name: {agent.name}\n"
-        return summary
\ No newline at end of file
+        return summary
+    
+Context.model_rebuild()
\ No newline at end of file