File size: 407 Bytes
3860419
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from typing import Any, List, Optional


class MockAI:
    def __init__(self, response: List):
        self.responses = iter(response)

    def start(self, system: str, user: Any, *, step_name: str) -> List[str]:
        return [next(self.responses)]

    def next(
        self, messages: List[str], prompt: Optional[str] = None, *, step_name: str
    ) -> List[str]:
        return [next(self.responses)]