Adding first set of tests for Suggesters

Wrote test classes and methods for the following classes:

SimpleModelSuggester
SimpleIdentificationSuggester
ModelSuggester
IdentificationSuggeste

Author: amit-sharma

pywhyllm/suggesters/identification_suggester.py

@@ -1,21 +1,19 @@
-from typing import List, Dict, Set, Tuple, Protocol
+from typing import List, Dict, Tuple
 from ..protocols import IdentifierProtocol
-from ..helpers import RelationshipStrategy, ModelType
 from .model_suggester import ModelSuggester
-from ..prompts import prompts as ps
 import guidance
 from guidance import system, user, assistant, gen
 import re
 
 
 class IdentificationSuggester(IdentifierProtocol):
-
     CONTEXT: str = """causal mechanisms"""
 
-    def __init__(self, llm):
-        if llm == 'gpt-4':
-            self.llm = guidance.models.OpenAI('gpt-4')
-            self.model_suggester = ModelSuggester('gpt-4')
+    def __init__(self, llm=None):
+        if llm is not None:
+            if (llm == 'gpt-4'):
+                self.llm = guidance.models.OpenAI('gpt-4')
+                self.model_suggester = ModelSuggester('gpt-4')
 
     # def suggest_estimand(
     #     self,
@@ -120,7 +118,7 @@ def suggest_backdoor(
             outcome: str,
             factors_list: list(),
             expertise_list: list(),
-            analysis_context: list() = CONTEXT,
+            analysis_context=CONTEXT,
             stakeholders: list() = None
     ):
         backdoor_set = self.model_suggester.suggest_confounders(
@@ -133,14 +131,14 @@ def suggest_backdoor(
         )
         return backdoor_set
 
-    #TODO:implement
+    # TODO:implement
     def suggest_frontdoor(
             self,
             treatment: str,
             outcome: str,
             factors_list: list(),
             expertise_list: list(),
-            analysis_context: list() = CONTEXT,
+            analysis_context=CONTEXT,
             stakeholders: list() = None
     ):
         pass
@@ -151,7 +149,7 @@ def suggest_mediators(
             outcome: str,
             factors_list: list(),
             expertise_list: list(),
-            analysis_context: list() = CONTEXT,
+            analysis_context=CONTEXT,
             stakeholders: list() = None
     ):
         expert_list: List[str] = list()
@@ -170,43 +168,28 @@ def suggest_mediators(
             if factors_list[i] != treatment and factors_list[i] != outcome:
                 edited_factors_list.append(factors_list[i])
 
-        if len(expert_list) > 1:
-            for expert in expert_list:
-                mediators_edges, mediators_list = self.request_mediators(
-                    treatment=treatment,
-                    outcome=outcome,
-                    analysis_context=analysis_context,
-                    domain_expertise=expert,
-                    factors_list=edited_factors_list,
-                    mediators_edges=mediators_edges
-                )
-                for m in mediators_list:
-                    if m not in mediators:
-                        mediators.append(m)
-        else:
+        for expert in expert_list:
             mediators_edges, mediators_list = self.request_mediators(
                 treatment=treatment,
                 outcome=outcome,
-                analysis_context=analysis_context,
-                domain_expertise=expert_list[0],
+                domain_expertise=expert,
                 factors_list=edited_factors_list,
                 mediators_edges=mediators_edges,
+                analysis_context=analysis_context
             )
-
             for m in mediators_list:
                 if m not in mediators:
                     mediators.append(m)
-
         return mediators_edges, mediators
 
     def request_mediators(
             self,
             treatment,
             outcome,
-            analysis_context,
             domain_expertise,
             factors_list,
-            mediators_edges
+            mediators_edges,
+            analysis_context=CONTEXT
     ):
         mediators: List[str] = list()
 
@@ -254,9 +237,7 @@ def request_mediators(
                         # to not add it twice into the list
                         if factor in factors_list and factor not in mediators:
                             mediators.append(factor)
-                    success = True
-                else:
-                    success = False
+                success = True
 
             except KeyError:
                 success = False
@@ -281,7 +262,7 @@ def suggest_ivs(
             outcome: str,
             factors_list: list(),
             expertise_list: list(),
-            analysis_context: list() = CONTEXT,
+            analysis_context=CONTEXT,
             stakeholders: list() = None
     ):
         expert_list: List[str] = list()
@@ -300,26 +281,20 @@ def suggest_ivs(
             if factors_list[i] != treatment and factors_list[i] != outcome:
                 edited_factors_list.append(factors_list[i])
 
-        if len(expert_list) > 1:
-            for expert in expert_list:
-                self.request_ivs(
-                    treatment=treatment,
-                    outcome=outcome,
-                    analysis_context=analysis_context,
-                    domain_expertise=expert,
-                    factors_list=edited_factors_list,
-                    iv_edges=iv_edges,
-                )
-        else:
-            self.request_ivs(
+        for expert in expert_list:
+            iv_edges, iv_list = self.request_ivs(
                 treatment=treatment,
                 outcome=outcome,
                 analysis_context=analysis_context,
-                domain_expertise=expert_list[0],
+                domain_expertise=expert,
                 factors_list=edited_factors_list,
                 iv_edges=iv_edges,
             )
 
+            for m in iv_list:
+                if m not in ivs:
+                    ivs.append(m)
+
         return iv_edges, ivs
 
     def request_ivs(
@@ -376,9 +351,7 @@ def request_ivs(
                     for factor in iv_factors:
                         if factor in factors_list and factor not in ivs:
                             ivs.append(factor)
-                    success = True
-                else:
-                    success = False
+                success = True
 
             except KeyError:
                 success = False
@@ -390,4 +363,4 @@ def request_ivs(
                 else:
                     iv_edges[(element, treatment)] = 1
 
-        return iv_edges
+        return iv_edges, ivs