Remove agent context from API

langstream-runtime/langstream-runtime-impl/src/main/python/langstream/api.py

@@ -16,20 +16,16 @@
 #
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Any, List, Tuple, Dict, Union, Optional
+from typing import Any, List, Tuple, Dict, Union
 
 __all__ = [
     "Record",
     "RecordType",
-    "AgentContext",
     "Agent",
     "Source",
     "Sink",
     "Processor",
     "CommitCallback",
-    "TopicConsumer",
-    "TopicProducer",
 ]
 
 
@@ -65,67 +61,6 @@ def headers(self) -> List[Tuple[str, Any]]:
 RecordType = Union[Record, list, tuple]
 
 
-class TopicConsumer(ABC):
-    """The topic consumer interface"""
-
-    def start(self):
-        """Start the consumer."""
-        pass
-
-    def close(self):
-        """Close the consumer"""
-        pass
-
-    def read(self) -> List[Record]:
-        """Read records from the topic."""
-        return []
-
-    def commit(self, records: List[Record]):
-        """Commit records."""
-        pass
-
-    def get_native_consumer(self) -> Any:
-        """Return the native wrapped consumer"""
-        pass
-
-    def get_info(self) -> Dict[str, Any]:
-        """Return the consumer info"""
-        return {}
-
-
-class TopicProducer(ABC):
-    """The topic producer interface"""
-
-    def start(self):
-        """Start the producer."""
-        pass
-
-    def close(self):
-        """Close the producer."""
-        pass
-
-    def write(self, records: List[Record]):
-        """Write records to the topic."""
-        pass
-
-    def get_native_producer(self) -> Any:
-        """Return the native wrapped producer"""
-        pass
-
-    def get_info(self) -> Dict[str, Any]:
-        """Return the producer info"""
-        return {}
-
-
-@dataclass
-class AgentContext(object):
-    """The Agent context"""
-
-    topic_consumer: Optional[TopicConsumer] = None
-    topic_producer: Optional[TopicProducer] = None
-    global_agent_id: Optional[str] = None
-
-
 class Agent(ABC):
     """The Agent interface"""
 
@@ -145,10 +80,6 @@ def agent_info(self) -> Dict[str, Any]:
         """Return the agent information."""
         return {}
 
-    def set_context(self, context: AgentContext):
-        """Set the agent context."""
-        pass
-
 
 class Source(Agent):
     """The Source agent interface

langstream-runtime/langstream-runtime-impl/src/main/python/langstream_runtime/api.py

@@ -16,13 +16,11 @@
 #
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from typing import Any, List, Tuple, Dict, Union, Optional
+from typing import Any, List, Tuple, Dict, Union
 
 __all__ = [
     "Record",
     "RecordType",
-    "AgentContext",
     "Agent",
     "Source",
     "Sink",
@@ -117,15 +115,6 @@ def get_info(self) -> Dict[str, Any]:
         return {}
 
 
-@dataclass
-class AgentContext(object):
-    """The Agent context"""
-
-    topic_consumer: Optional[TopicConsumer] = None
-    topic_producer: Optional[TopicProducer] = None
-    global_agent_id: Optional[str] = None
-
-
 class Agent(ABC):
     """The Agent interface"""
 
@@ -145,10 +134,6 @@ def agent_info(self) -> Dict[str, Any]:
         """Return the agent information."""
         return {}
 
-    def set_context(self, context: AgentContext):
-        """Set the agent context."""
-        pass
-
 
 class Source(Agent):
     """The Source agent interface

langstream-runtime/langstream-runtime-impl/src/main/python/langstream_runtime/runtime.py

@@ -32,7 +32,6 @@
     Processor,
     Record,
     CommitCallback,
-    AgentContext,
     RecordType,
 )
 from .util import SingleRecordProcessor, SimpleRecord
@@ -92,7 +91,7 @@ class ComponentType(str, Enum):
 class RuntimeAgent(Agent):
     def __init__(
         self,
-        agent: Agent,
+        agent: Union[Source, Sink, Processor],
         component_type: ComponentType,
         agent_id,
         agent_type,
@@ -136,9 +135,6 @@ def get_agent_status(self) -> List[Dict[str, Any]]:
             }
         ]
 
-    def set_context(self, context: AgentContext):
-        call_method_if_exists(self.agent, "set_context", context)
-
 
 class RuntimeSource(RuntimeAgent, Source):
     def __init__(self, source: Source, agent_id, agent_type, started_at=None):
@@ -329,15 +325,6 @@ def run(configuration, agent=None, agent_info: AgentInfo = AgentInfo(), max_loop
 
     agent_info.processor = processor
 
-    agent_context = AgentContext(
-        topic_consumer=consumer,
-        topic_producer=producer,
-        global_agent_id=application_agent_id,
-    )
-
-    for component in {a.agent: a for a in {source, sink, processor}}.values():
-        component.set_context(agent_context)
-
     run_main_loop(
         source,
         sink,