agents/tools/smart_domain/association_impl.py

from langchain import LLMChain, PromptTemplate
from agents.tools.smart_domain.common import getPrefix
from models import llm
from langchain.agents import tool

association_impl_architecture = """* Association Impletation: This component provide implementation of Association of domain layer, rely on Repository. 
---eaxmple code: 
    @Component
    @RequiredArgsConstructor
    public class FeaturesImpl implements Features{{
        private final FeatureDbRepository featureDbRepository;

        Flux<Feature> findAll() {{
            return featureDbRepository.findAll().map(FeatureDb::toFeature);
        }}

        Mono<Feature> save(Feature feature) {{
            return featureDbRepository.save(FeatureDb.fromFeature(feature)).map(FeatureDb::toFeature);
        }}
    }}
---end of eaxmple code"""

association_impl_test_strategy="""For Association Impletation，we writ unit test, and stub repository method with Mockito.
    ---eaxmple code:
        @ExtendWith(MockitoExtension.class)
        class FeatureImplTest {{
            @Mock
            FeatureDbRepository repository;

            Features features;

            @BeforeEach
            void setUp() {{
                features = new FeaturesImpl(repository);
            }}

            @Test
            void should_add_success() {{
                when(repository.save(any(FeatureDb.class))).thenAnswer(invocation -> {{
                    FeatureDb featureDb = invocation.getArgument(0);
                    return Mono.just(featureDb);
                }});

                features.add(createFeature("featureKey1"))
                        .as(StepVerifier::create)
                        .expectNextMatches(config -> config.getId().featureKey().equals("featureKey1")
                                && config.getDescription().updatedAt() != null
                                && config.getDescription().createdAt() != null
                        )
                        .verifyComplete();
            }}

            @Test
            void should_add_return_error_when_repository_save_error() {{
                Feature feature = createFeature("featureKey1");

                when(repository.save(any(FeatureDb.class))).thenReturn(Mono.error(new DuplicateKeyException("save error")));

                features.add(feature)
                        .as(StepVerifier::create)
                        .expectError()
                        .verify();
            }}
        }}"""

association_impl_task = """Your task is to generate the Association Impletation tests and product code."""

association_impl_teck_stack = """Java17、reactor、lombok、Junit5、reactor test、Mockito、 Spring Data Reactive Couchbase"""

ASSOCIATION_IMPL = getPrefix(association_impl_task, association_impl_teck_stack, association_impl_architecture, association_impl_test_strategy) + """

Use the following format:
request: the request that you need to fulfill

Association Impletation:
```
the Association Impletation code that you write to fulfill the request, follow TechStack and Architecture
```

Test:
```
the test code of Association Impletation that you write to fulfill the request, follow TechStack Architecture and TestStrategy
```

request: {input}
"""

ASSOCIATION_IMPL_PROMOPT = PromptTemplate(input_variables=["input"], template=ASSOCIATION_IMPL,)

asociationChain = LLMChain(llm = llm(temperature=0.1), prompt=ASSOCIATION_IMPL_PROMOPT)


@tool("Generate Asociation Impletation Code", return_direct=True)
def asociationImplCodeGenerator(input: str) -> str:
    '''useful for when you need to generate Asociation Impletation code'''
    response = asociationChain.run(input)
    return response